from __future__ import annotations

import re
from dataclasses import dataclass
from typing import TYPE_CHECKING

from kernbench.common.types import Completion, Trace

if TYPE_CHECKING:
    from kernbench.sim_engine.engine import GraphEngine


@dataclass(frozen=True)
class BenchResult:
    completion: Completion
    correlation_id: str
    trace: Trace | None = None
    traces: list[dict] | None = None
    engine: GraphEngine | None = None

    def summary_text(self) -> str:
        if self.completion.ok:
            return f"[OK] correlation_id={self.correlation_id}"
        code = self.completion.error_code or "ERROR"
        msg = self.completion.error_message or ""
        return f"[FAIL:{code}] correlation_id={self.correlation_id} {msg}".rstrip()


@dataclass(frozen=True)
class DeviceSelector:
    """
    Device selector.

    Supported:
      - "all"      : all SIPs in the tray topology
      - "sip:<N>"  : a single SIP index
    """

    raw: str  # "all" or "sip:<N>"

    @property
    def is_all(self) -> bool:
        return self.raw == "all"

    @property
    def sip_index(self) -> int:
        if self.is_all:
            raise ValueError("DeviceSelector is 'all'; no single sip_index.")
        m = re.fullmatch(r"sip:?(\d+)", self.raw)
        if not m:
            raise ValueError(
                f"Invalid device '{self.raw}'. Expected 'all' or 'sip:<N>' (e.g., sip:0)."
            )
        return int(m.group(1))


def resolve_device(raw: str | None) -> DeviceSelector:
    """
    Resolve the CLI --device string into a DeviceSelector.

    Semantics:
      - if omitted/empty -> "all"
      - else accept "all" or "sip:<N>"
    """
    if raw is None or raw.strip() == "":
        return DeviceSelector(raw="all")

    raw = raw.strip().lower()
    if raw == "all":
        return DeviceSelector(raw="all")

    m = re.fullmatch(r"sip:?(\d+)", raw)
    if not m:
        raise ValueError(f"Invalid device '{raw}'. Expected 'all' or 'sip:<N>' (e.g., sip:0).")
    raw = f"sip:{m.group(1)}"  # normalize to sip:N format

    return DeviceSelector(raw=raw)
