from __future__ import annotations

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


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


def parse_iso_datetime(value: str | None) -> datetime | None:
    if not value:
        return None
    text = str(value).strip()
    if not text:
        return None
    if text.endswith("Z"):
        text = text[:-1] + "+00:00"
    try:
        dt = datetime.fromisoformat(text)
    except ValueError:
        return None
    if dt.tzinfo is None:
        return dt.replace(tzinfo=timezone.utc)
    return dt.astimezone(timezone.utc)


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 filter_recent(records: list[dict[str, Any]], *, days: int | None) -> list[dict[str, Any]]:
    if not days or days <= 0:
        return records
    threshold = datetime.now(timezone.utc) - timedelta(days=days)
    result = []
    for record in records:
        dt = parse_iso_datetime(record.get("timestamp") or record.get("logged_at"))
        if dt is None or dt >= threshold:
            result.append(record)
    return result


def week_bucket(dt: datetime | None) -> str:
    if dt is None:
        return "unknown"
    iso_year, iso_week, _ = dt.isocalendar()
    return f"{iso_year}-W{iso_week:02d}"


def top_counter(counter: Counter[str], *, limit: int) -> list[dict[str, Any]]:
    return [{"key": key, "count": count} for key, count in counter.most_common(limit)]


def summarize(records: list[dict[str, Any]], *, limit: int = 10) -> dict[str, Any]:
    weekly_turns: dict[str, dict[str, int]] = defaultdict(lambda: {"turns": 0, "success": 0, "partial": 0, "failure": 0})
    failure_types = Counter()
    repeated_failed_prompts = Counter()
    model_breakdown: dict[str, dict[str, Any]] = defaultdict(lambda: {"calls": 0, "_latencies": [], "avg_latency_ms": None})
    tool_breakdown: dict[str, dict[str, Any]] = defaultdict(lambda: {"calls": 0, "_durations": [], "avg_duration_ms": None})
    rag_breakdown: dict[str, dict[str, int]] = defaultdict(lambda: {"retrievals": 0, "sources": 0})

    for record in records:
        event_type = str(record.get("event_type") or "")
        if event_type == "turn_audit":
            dt = parse_iso_datetime(record.get("timestamp"))
            week = week_bucket(dt)
            outcome = str(record.get("outcome") or "")
            weekly_turns[week]["turns"] += 1
            if outcome in {"success", "partial", "failure"}:
                weekly_turns[week][outcome] += 1
            failure_type = str(record.get("failure_type") or "")
            if failure_type and failure_type != "none":
                failure_types[failure_type] += 1
                prompt = " ".join(str(record.get("user_text") or "").split())
                if prompt:
                    repeated_failed_prompts[prompt] += 1
        elif event_type == "model_calls":
            key = f"{record.get('route_name') or ''} :: {record.get('served_model_name') or ''}"
            model_breakdown[key]["calls"] += 1
            latency = record.get("latency_ms")
            if isinstance(latency, int | float):
                model_breakdown[key]["_latencies"].append(float(latency))
        elif event_type == "tool_calls":
            key = str(record.get("tool_name") or "")
            tool_breakdown[key]["calls"] += 1
            duration = record.get("duration_ms")
            if isinstance(duration, int | float):
                tool_breakdown[key]["_durations"].append(float(duration))
        elif event_type == "retrieval_events":
            key = f"{record.get('plugin') or ''} :: {record.get('profile') or ''}"
            rag_breakdown[key]["retrievals"] += 1
            sources = record.get("sources_json") or []
            if isinstance(sources, list):
                rag_breakdown[key]["sources"] += len(sources)

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

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

    rag_rows = []
    for key, item in sorted(rag_breakdown.items()):
        sources = item["sources"]
        retrievals = item["retrievals"]
        rag_rows.append(
            {
                "key": key,
                "retrievals": retrievals,
                "avg_sources": round(sources / retrievals, 2) if retrievals else None,
            }
        )

    weekly_rows = [{"week": week, **stats} for week, stats in sorted(weekly_turns.items())]

    return {
        "records": len(records),
        "weekly_turns": weekly_rows,
        "top_failure_types": top_counter(failure_types, limit=limit),
        "top_repeated_failed_prompts": top_counter(repeated_failed_prompts, limit=limit),
        "model_breakdown": model_rows[:limit],
        "tool_breakdown": tool_rows[:limit],
        "rag_breakdown": rag_rows[:limit],
    }


def write_report(path: Path, payload: dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(payload, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")


def main() -> int:
    parser = argparse.ArgumentParser(description="Build a quality report from observability logs.")
    parser.add_argument("--report-dir", type=Path, default=DEFAULT_REPORT_DIR)
    parser.add_argument("--days", type=int, default=7, help="Include only the latest N days. Use 0 to include all records.")
    parser.add_argument("--limit", type=int, default=10, help="Top-N rows per section.")
    parser.add_argument("--output", type=Path, help="Optional JSON output path.")
    args = parser.parse_args()

    records = iter_records(args.report_dir)
    filtered = filter_recent(records, days=args.days)
    payload = summarize(filtered, limit=args.limit)
    print(json.dumps(payload, ensure_ascii=False, indent=2))
    if args.output:
        write_report(args.output, payload)
    return 0


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