"""Parse agent_trace.log into structured events."""

from __future__ import annotations

import re
from dataclasses import dataclass, field
from pathlib import Path

# Matches: [13:34:11 PST / 01:34:11 +04 / 21:34:11 UTC]
_TIMESTAMP_RE = re.compile(r"^\[(.+?UTC)\]\s+")

# START appendix_tester (7ebb3c7a) (call) model=haiku
_START_RE = re.compile(
    r"^START\s+(\S+)"
    r"(?:\s+\(([0-9a-f]+)\))?"
    r"\s+\((call|resume)\)"
    r"\s+model=(\S+)"
)

# END appendix_tester (7ebb3c7a) session=xxx | duration=4.7s | cost=$0.0052 | status=DONE
_END_RE = re.compile(
    r"^END\s+(\S+)"
    r"(?:\s+\(([0-9a-f]+)\))?"
    r"\s+(.*)"
)

# ERROR agent_name (call_id) error message
_ERROR_RE = re.compile(
    r"^ERROR\s+(\S+)"
    r"(?:\s+\(([0-9a-f]+)\))?"
    r"(?:\s+(.*))?"
)


def _parse_duration(raw: str) -> float | None:
    """Parse duration strings like '4.7s', '2m 13s', or '2m' into seconds."""
    if not raw:
        return None
    m = re.match(r"^(\d+)m\s+(\d+)s$", raw)
    if m:
        return int(m.group(1)) * 60 + int(m.group(2))
    m = re.match(r"^(\d+)m$", raw)
    if m:
        return int(m.group(1)) * 60
    m = re.match(r"^([\d.]+)s$", raw)
    if m:
        return float(m.group(1))
    return None


def _parse_cost(raw: str) -> float | None:
    """Parse cost strings like '$0.0052' into float."""
    if not raw or raw == "n/a":
        return None
    if raw.startswith("$"):
        try:
            return float(raw[1:])
        except ValueError:
            return None
    return None


@dataclass
class TraceEvent:
    timestamp_utc: str = ""
    event_type: str = ""       # "start", "end", "error"
    agent_name: str = ""
    call_id: str = ""
    model: str = ""
    mode: str = ""             # "call" or "resume"
    duration: str = ""         # raw string e.g. "4.7s"
    cost: str = ""             # raw string e.g. "$0.0052"
    status: str = ""
    error_message: str = ""

    @property
    def duration_seconds(self) -> float | None:
        return _parse_duration(self.duration)

    @property
    def cost_usd(self) -> float | None:
        return _parse_cost(self.cost)


@dataclass
class TraceSummary:
    total_calls: int = 0
    total_duration_seconds: float = 0.0
    total_cost_usd: float = 0.0
    events: list[TraceEvent] = field(default_factory=list)
    agents: dict[str, int] = field(default_factory=dict)


def parse_trace_line(line: str) -> TraceEvent | None:
    """Parse a single trace log line into a TraceEvent, or None if unparseable."""
    line = line.strip()
    if not line:
        return None

    ts_match = _TIMESTAMP_RE.match(line)
    if not ts_match:
        return None

    timestamp = ts_match.group(1)
    rest = line[ts_match.end():]

    # Try START
    m = _START_RE.match(rest)
    if m:
        return TraceEvent(
            timestamp_utc=timestamp,
            event_type="start",
            agent_name=m.group(1),
            call_id=m.group(2) or "",
            mode=m.group(3),
            model=m.group(4),
        )

    # Try END
    m = _END_RE.match(rest)
    if m:
        agent_name = m.group(1)
        call_id = m.group(2) or ""
        info_str = m.group(3) or ""
        # Parse pipe-separated key=value pairs (values may contain spaces, e.g. "2m 19s")
        kvs = {}
        for segment in info_str.split(" | "):
            segment = segment.strip()
            eq = segment.find("=")
            if eq > 0:
                kvs[segment[:eq]] = segment[eq + 1:]
        return TraceEvent(
            timestamp_utc=timestamp,
            event_type="end",
            agent_name=agent_name,
            call_id=call_id,
            duration=kvs.get("duration", ""),
            cost=kvs.get("cost", ""),
            status=kvs.get("status", ""),
        )

    # Try ERROR
    m = _ERROR_RE.match(rest)
    if m:
        return TraceEvent(
            timestamp_utc=timestamp,
            event_type="error",
            agent_name=m.group(1),
            call_id=m.group(2) or "",
            error_message=m.group(3) or "",
        )

    return None


def parse_trace_log(path: Path | str) -> TraceSummary:
    """Parse a full agent_trace.log file into a TraceSummary."""
    path = Path(path)
    summary = TraceSummary()

    if not path.is_file():
        return summary

    text = path.read_text()
    for line in text.splitlines():
        event = parse_trace_line(line)
        if event is None:
            continue
        summary.events.append(event)

        if event.event_type == "end":
            summary.total_calls += 1
            summary.agents[event.agent_name] = summary.agents.get(event.agent_name, 0) + 1
            if event.duration_seconds is not None:
                summary.total_duration_seconds += event.duration_seconds
            if event.cost_usd is not None:
                summary.total_cost_usd += event.cost_usd

    return summary
