"""Parse agent run logs (.jsonl or legacy .log) into structured events."""

from __future__ import annotations

import json
import re
from pathlib import Path
from typing import Any


_TITLE_RE = re.compile(
    r"^(RUN START|RUN END|LLM INPUT|LLM OUTPUT|TOOL CALL|SKILL LOADED)"
    r"(?:\s*\|\s*(.+))?$",
)


def _title_to_event(title: str) -> tuple[str, dict[str, Any]]:
    """Map a human log title to (event_type, extra_fields)."""
    m = _TITLE_RE.match(title.strip())
    if not m:
        return "unknown", {"title": title}

    kind = m.group(1)
    rest = m.group(2) or ""
    extra: dict[str, Any] = {}

    iter_m = re.search(r"iteration=(\d+)", rest)
    if iter_m:
        extra["iteration"] = int(iter_m.group(1))

    mapping = {
        "RUN START": "run_start",
        "RUN END": "run_end",
        "LLM INPUT": "llm_input",
        "LLM OUTPUT": "llm_output",
        "TOOL CALL": "tool_call",
        "SKILL LOADED": "skill_loaded",
    }
    return mapping.get(kind, "unknown"), extra


def parse_jsonl(path: Path) -> list[dict[str, Any]]:
    """Parse structured JSONL event stream."""
    events: list[dict[str, Any]] = []
    with path.open(encoding="utf-8") as fh:
        for line_no, line in enumerate(fh, 1):
            line = line.strip()
            if not line:
                continue
            try:
                events.append(json.loads(line))
            except json.JSONDecodeError as exc:
                raise ValueError(f"Invalid JSONL at {path}:{line_no}: {exc}") from exc
    return events


def parse_legacy_log(path: Path) -> list[dict[str, Any]]:
    """Parse human-readable .log blocks into JSONL-compatible events."""
    text = path.read_text(encoding="utf-8")
    # Split on separator lines that follow a timestamp prefix
    parts = re.split(
        r"\[\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2}\] =+\n",
        text,
    )
    events: list[dict[str, Any]] = []
    seq = 0

    for part in parts:
        part = part.strip()
        if not part:
            continue
        # Format: TITLE\n====\n{json}
        lines = part.split("\n", 2)
        if len(lines) < 3:
            continue
        title = lines[0].strip()
        # lines[1] is separator
        body = lines[2]
        # Strip trailing separator residue if any
        body = re.sub(r"\n=+\s*$", "", body).strip()
        try:
            data = json.loads(body)
        except json.JSONDecodeError:
            continue

        event_type, extra = _title_to_event(title)
        seq += 1

        # Normalize legacy llm_output shape: reasoning lived only under parsed/raw
        if event_type == "llm_output" and "content" not in data:
            parsed = data.get("parsed") or {}
            data = {
                **data,
                "content": parsed.get("content"),
                "reasoning": parsed.get("reasoning")
                or ((data.get("raw") or {}).get("choices") or [{}])[0]
                .get("message", {})
                .get("reasoning"),
                "tool_calls": parsed.get("tool_calls") or [],
                "usage": (data.get("raw") or {}).get("usage"),
            }

        if event_type == "tool_call" and "arguments_parsed" not in data:
            args = data.get("arguments", "")
            try:
                data["arguments_parsed"] = json.loads(args) if isinstance(args, str) else args
            except (json.JSONDecodeError, TypeError):
                data["arguments_parsed"] = args
            result = data.get("result", "")
            try:
                data["result_parsed"] = json.loads(result) if isinstance(result, str) else result
            except (json.JSONDecodeError, TypeError):
                data["result_parsed"] = result

        record: dict[str, Any] = {
            "event": event_type,
            "seq": seq,
            "run_id": data.get("run_id"),
            "data": data,
        }
        if "iteration" in extra:
            record["iteration"] = extra["iteration"]
        elif "iteration" in data:
            record["iteration"] = data["iteration"]
        events.append(record)

    return events


def load_run_events(path: Path | str) -> list[dict[str, Any]]:
    """
    Load events from a run file.

    Accepts ``.jsonl``, ``.log``, or a stem without extension.
    Prefers ``.jsonl`` when both exist.
    """
    path = Path(path)

    if path.is_dir():
        raise ValueError(f"Expected a run file, got directory: {path}")

    candidates: list[Path] = []
    if path.suffix == ".jsonl":
        candidates = [path]
    elif path.suffix == ".log":
        jsonl = path.with_suffix(".jsonl")
        candidates = [jsonl, path] if jsonl.exists() else [path]
    elif path.exists():
        candidates = [path]
    else:
        # Stem lookup: run_xxx → try .jsonl then .log
        candidates = [path.with_suffix(".jsonl"), path.with_suffix(".log"), Path(str(path) + ".jsonl"), Path(str(path) + ".log")]

    for candidate in candidates:
        if not candidate.exists():
            continue
        if candidate.suffix == ".jsonl":
            return parse_jsonl(candidate)
        if candidate.suffix == ".log":
            return parse_legacy_log(candidate)
        # Unknown suffix: try jsonl lines first
        try:
            return parse_jsonl(candidate)
        except ValueError:
            return parse_legacy_log(candidate)

    raise FileNotFoundError(f"No run log found for: {path}")


def summarize_run(events: list[dict[str, Any]]) -> dict[str, Any]:
    """Extract run-level metadata from event list."""
    start = next((e for e in events if e.get("event") == "run_start"), None)
    end = next((e for e in events if e.get("event") == "run_end"), None)
    start_data = (start or {}).get("data") or {}
    end_data = (end or {}).get("data") or {}

    iterations = {
        e.get("iteration")
        for e in events
        if e.get("iteration") is not None
    }
    tool_events = [e for e in events if e.get("event") == "tool_call"]

    return {
        "run_id": start_data.get("run_id") or end_data.get("run_id") or (start or {}).get("run_id"),
        "model": start_data.get("model"),
        "user_input": start_data.get("user_input"),
        "iterations": end_data.get("iterations") or (max(iterations) if iterations else 0),
        "tool_calls_made": end_data.get("tool_calls_made", len(tool_events)),
        "skills_used": end_data.get("skills_used") or [],
        "final_content": end_data.get("final_content"),
        "event_count": len(events),
    }