from __future__ import annotations

import argparse
import json
from collections import Counter, defaultdict
from pathlib import Path
from typing import Any


ROOT = Path(__file__).resolve().parents[1]
DEFAULT_REPORT_DIR = ROOT / "reports" / "observability"


def iter_records(report_dir: Path) -> list[dict[str, Any]]:
    records: list[dict[str, Any]] = []
    if not report_dir.exists():
        return records
    for path in sorted(report_dir.rglob("*.jsonl")):
        with path.open("r", encoding="utf-8") as handle:
            for line_number, line in enumerate(handle, start=1):
                line = line.strip()
                if not line:
                    continue
                try:
                    record = json.loads(line)
                except json.JSONDecodeError as exc:
                    raise ValueError(f"{path}:{line_number}: invalid JSONL: {exc}") from exc
                if isinstance(record, dict):
                    records.append(record)
    return records


def summarize(records: list[dict[str, Any]]) -> dict[str, Any]:
    by_event = Counter()
    by_status = Counter()
    failure_types = Counter()
    plugin_turns = defaultdict(lambda: {"turns": 0, "success": 0, "partial": 0, "failure": 0})
    model_routes = defaultdict(lambda: {"calls": 0, "avg_latency_ms": None, "_latencies": []})
    tool_summary = defaultdict(lambda: {"calls": 0, "avg_duration_ms": None, "_durations": []})

    for record in records:
        event_type = str(record.get("event_type") or "")
        by_event[event_type] += 1

        if event_type == "access_events":
            by_status[str(record.get("status_code") or "")] += 1
        elif event_type == "turn_audit":
            plugin = str(record.get("plugin") or "")
            outcome = str(record.get("outcome") or "")
            failure_type = str(record.get("failure_type") or "")
            item = plugin_turns[plugin]
            item["turns"] += 1
            if outcome in {"success", "partial", "failure"}:
                item[outcome] += 1
            if failure_type and failure_type != "none":
                failure_types[failure_type] += 1
        elif event_type == "model_calls":
            route = str(record.get("route_name") or "")
            item = model_routes[route]
            item["calls"] += 1
            latency = record.get("latency_ms")
            if isinstance(latency, int | float):
                item["_latencies"].append(float(latency))
        elif event_type == "tool_calls":
            tool_name = str(record.get("tool_name") or "")
            item = tool_summary[tool_name]
            item["calls"] += 1
            duration = record.get("duration_ms")
            if isinstance(duration, int | float):
                item["_durations"].append(float(duration))

    model_route_rows = []
    for route_name, item in sorted(model_routes.items()):
        latencies = item.pop("_latencies")
        item["route_name"] = route_name
        item["avg_latency_ms"] = round(sum(latencies) / len(latencies)) if latencies else None
        model_route_rows.append(item)

    tool_rows = []
    for tool_name, item in sorted(tool_summary.items()):
        durations = item.pop("_durations")
        item["tool_name"] = tool_name
        item["avg_duration_ms"] = round(sum(durations) / len(durations)) if durations else None
        tool_rows.append(item)

    return {
        "records": len(records),
        "by_event_type": dict(sorted(by_event.items())),
        "http_status_counts": dict(sorted(by_status.items())),
        "failure_type_counts": dict(sorted(failure_types.items())),
        "turns_by_plugin": {key: value for key, value in sorted(plugin_turns.items())},
        "model_routes": model_route_rows,
        "tool_calls": tool_rows,
    }


def main() -> int:
    parser = argparse.ArgumentParser(description="Summarize observability JSONL reports.")
    parser.add_argument("--report-dir", type=Path, default=DEFAULT_REPORT_DIR)
    args = parser.parse_args()

    records = iter_records(args.report_dir)
    print(json.dumps(summarize(records), ensure_ascii=False, indent=2))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())