"""Standardized LLM Telemetry Emitter.

Constitutional mandate: every LLM invocation across the ABI fleet MUST emit
a standardized telemetry record, regardless of the agent's role.

Usage::

    emitter = TelemetryEmitter(output_dir=run_dir, run_id="dogfood-20260307T171528")
    emitter.emit(
        provider="anthropic", model="claude-opus-4-6", role="builder",
        duration_ms=42000, task_id="IGF-001", repo="abi-orchestrator",
    )
    print(emitter.summary_line())
"""

from __future__ import annotations

import json
import uuid
from datetime import UTC, datetime
from pathlib import Path
from typing import Any

VALID_ROLES = frozenset(
    {
        "builder",
        "verifier",
        "auditor",
        "planner",
        "evaluator",
        "reviewer",
        "classifier",
        "router",
        "extractor",
        "summarizer",
        "orchestrator",   # system-level events (task lifecycle)
        "task_summary",   # per-task completion aggregates
    }
)

VALID_PROVIDERS = frozenset(
    {
        "anthropic",
        "openai",
        "google",
        "ollama",
        "replay",
        "mock",
        "system",  # internal orchestrator events (no LLM call)
    }
)

VALID_STATUSES = frozenset(
    {
        "success",
        "failure",
        "timeout",
        "rate_limited",
        "started",  # task lifecycle: execution began
    }
)


class TelemetryEmitter:
    """Emit and aggregate LLM telemetry records to NDJSON.

    Each record is appended to ``llm_telemetry.ndjson`` in the output
    directory immediately on :meth:`emit`, ensuring no data is lost
    even if the process crashes.
    """

    def __init__(self, output_dir: Path, run_id: str) -> None:
        self.output_dir = Path(output_dir)
        self.run_id = run_id
        self._file = self.output_dir / "llm_telemetry.ndjson"
        self._records: list[dict[str, Any]] = []
        self.output_dir.mkdir(parents=True, exist_ok=True)

    def emit(
        self,
        *,
        provider: str,
        model: str,
        role: str,
        duration_ms: int,
        status: str = "success",
        task_id: str | None = None,
        task_category: str | None = None,
        repo: str | None = None,
        tokens_in: int | None = None,
        tokens_out: int | None = None,
        cost_usd: float | None = None,
        error_class: str | None = None,
    ) -> dict[str, Any]:
        """Emit a single telemetry record.

        Validates role, provider, and status against allowed enums.
        Appends to the NDJSON file and returns the record dict.
        """
        if role not in VALID_ROLES:
            raise ValueError(f"Invalid role '{role}'. Must be one of: {sorted(VALID_ROLES)}")
        if provider not in VALID_PROVIDERS:
            raise ValueError(
                f"Invalid provider '{provider}'. Must be one of: {sorted(VALID_PROVIDERS)}"
            )
        if status not in VALID_STATUSES:
            raise ValueError(
                f"Invalid status '{status}'. Must be one of: {sorted(VALID_STATUSES)}"
            )

        total_tokens: int | None = None
        if tokens_in is not None or tokens_out is not None:
            total_tokens = (tokens_in or 0) + (tokens_out or 0)

        record: dict[str, Any] = {
            "invocation_id": str(uuid.uuid4()),
            "run_id": self.run_id,
            "timestamp": datetime.now(UTC).isoformat(),
            "provider": provider,
            "model": model,
            "role": role,
            "task_id": task_id,
            "task_category": task_category,
            "repo": repo,
            "tokens_in": tokens_in,
            "tokens_out": tokens_out,
            "total_tokens": total_tokens,
            "cost_usd": cost_usd,
            "duration_ms": duration_ms,
            "status": status,
            "error_class": error_class,
        }

        self._records.append(record)

        with open(self._file, "a", encoding="utf-8") as f:
            f.write(json.dumps(record, default=str) + "\n")

        return record

    def summary(self) -> dict[str, Any]:
        """Aggregate telemetry into a summary dict."""
        total_tokens = sum(r["total_tokens"] or 0 for r in self._records)
        total_cost = sum(r["cost_usd"] or 0.0 for r in self._records)
        total_calls = len(self._records)
        total_duration = sum(r["duration_ms"] for r in self._records)

        by_role: dict[str, dict[str, Any]] = {}
        by_provider: dict[str, dict[str, Any]] = {}
        by_model: dict[str, dict[str, Any]] = {}

        for r in self._records:
            for key, bucket in [("role", by_role), ("provider", by_provider), ("model", by_model)]:
                v = r[key]
                if v not in bucket:
                    bucket[v] = {"calls": 0, "tokens": 0, "cost": 0.0, "duration_ms": 0}
                bucket[v]["calls"] += 1
                bucket[v]["tokens"] += r["total_tokens"] or 0
                bucket[v]["cost"] += r["cost_usd"] or 0.0
                bucket[v]["duration_ms"] += r["duration_ms"]

        success_count = sum(1 for r in self._records if r["status"] == "success")

        return {
            "total_calls": total_calls,
            "total_tokens": total_tokens,
            "total_cost_usd": round(total_cost, 4),
            "total_duration_ms": total_duration,
            "by_role": by_role,
            "by_provider": by_provider,
            "by_model": by_model,
            "success_rate": success_count / max(total_calls, 1),
        }

    def summary_line(self) -> str:
        """One-line summary suitable for console output."""
        s = self.summary()
        tokens = f"{s['total_tokens']:,}" if s["total_tokens"] else "N/A"
        cost = f"${s['total_cost_usd']:.2f}" if s["total_cost_usd"] else "N/A"
        return (
            f"LLM Telemetry: {s['total_calls']} calls | "
            f"{tokens} tokens | {cost} | "
            f"{s['success_rate']:.0%} success"
        )
