# module12/observability.py
"""Tracing, structured logs, dashboards, alerting on quality, and replaying a failure locally."""
from __future__ import annotations

import json
import time
import uuid
from contextlib import contextmanager
from dataclasses import asdict, dataclass, field
from datetime import datetime, timezone
from pathlib import Path

import numpy as np


@dataclass
class Span:
    name: str
    start: float
    seconds: float = 0.0
    attributes: dict = field(default_factory=dict)
    error: str | None = None


@dataclass
class Trace:
    """One request, end to end, under a single correlation id."""
    correlation_id: str = field(default_factory=lambda: uuid.uuid4().hex[:12])
    query: str = ""
    tenant: str = ""
    user: str = ""
    started: str = field(default_factory=lambda: datetime.now(timezone.utc).isoformat(timespec="seconds"))
    spans: list[Span] = field(default_factory=list)
    attributes: dict = field(default_factory=dict)

    @contextmanager
    def span(self, name: str, **attributes):
        span = Span(name=name, start=time.perf_counter(), attributes=dict(attributes))
        self.spans.append(span)
        try:
            yield span
        except Exception as err:                      # record the failure, then let it propagate
            span.error = f"{type(err).__name__}: {err}"
            raise
        finally:
            span.seconds = time.perf_counter() - span.start

    def set(self, **attributes) -> None:
        self.attributes.update(attributes)

    @property
    def total_seconds(self) -> float:
        return sum(s.seconds for s in self.spans)

    def breakdown(self) -> dict:
        return {s.name: round(s.seconds * 1000, 1) for s in self.spans}

    def to_record(self) -> dict:
        return {"correlation_id": self.correlation_id, "started": self.started, "query": self.query,
                "tenant": self.tenant, "user": self.user,
                "total_ms": round(self.total_seconds * 1000, 1),
                "spans": [{**asdict(s), "start": None, "seconds": round(s.seconds * 1000, 1)}
                          for s in self.spans],
                **self.attributes}


class TraceLog:
    """Structured logs, one JSON object per line. The input to dashboards and to replay."""

    def __init__(self, path: str = "data/traces.jsonl", redactor=None):
        self.path = Path(path)
        self.path.parent.mkdir(parents=True, exist_ok=True)
        self.redactor = redactor                      # e.g. module12.governance.redact

    def write(self, trace: Trace) -> dict:
        record = trace.to_record()
        if self.redactor:
            record["query"] = self.redactor(record["query"])
            record["user"] = self.redactor(record.get("user", ""))
        with self.path.open("a", encoding="utf-8") as f:
            f.write(json.dumps(record, default=str) + "\n")
        return record

    def read(self, limit: int | None = None) -> list[dict]:
        if not self.path.exists():
            return []
        rows = [json.loads(line) for line in self.path.read_text(encoding="utf-8").splitlines()]
        return rows[-limit:] if limit else rows

    def find(self, correlation_id: str) -> dict | None:
        return next((r for r in self.read() if r["correlation_id"] == correlation_id), None)


# ---------------------------------------------------------------- dashboards
def percentile(values: list[float], p: float) -> float:
    return float(np.percentile(values, p)) if values else 0.0


def dashboard(records: list[dict]) -> dict:
    """The numbers to put on a wall: latency, quality, spend, and where the time goes."""
    if not records:
        return {}
    latencies = [r["total_ms"] for r in records]
    stage_totals: dict[str, list[float]] = {}
    for record in records:
        for span in record["spans"]:
            stage_totals.setdefault(span["name"], []).append(span["seconds"])
    return {
        "requests": len(records),
        "latency_ms": {"p50": round(percentile(latencies, 50), 1),
                       "p95": round(percentile(latencies, 95), 1),
                       "p99": round(percentile(latencies, 99), 1)},
        "stage_p95_ms": {name: round(percentile(values, 95), 1)
                         for name, values in sorted(stage_totals.items())},
        "retrieval_hit_rate": round(float(np.mean([r.get("retrieval_hit", 0) for r in records])), 3),
        "abstention_rate": round(float(np.mean([r.get("refused", False) for r in records])), 3),
        "cache_hit_rate": round(float(np.mean([bool(r.get("cache_hit")) for r in records])), 3),
        "error_rate": round(float(np.mean([bool(r.get("error")) for r in records])), 3),
        "degraded_rate": round(float(np.mean([bool(r.get("degraded")) for r in records])), 3),
        "spend_usd": round(sum(r.get("cost_usd", 0.0) for r in records), 4),
        "cost_per_answered_question": round(
            sum(r.get("cost_usd", 0.0) for r in records)
            / max(1, sum(not r.get("refused", False) for r in records)), 5),
    }


# ---------------------------------------------------------------- alerting
@dataclass
class AlertRule:
    """Alert on quality, not only on errors: a system that answers wrongly returns HTTP 200."""
    name: str
    metric: str                       # dotted path into the dashboard
    minimum: float | None = None
    maximum: float | None = None
    window: int = 100                 # how many recent requests to look at

    def value(self, dash: dict):
        value = dash
        for part in self.metric.split("."):
            value = value[part]
        return value

    def check(self, dash: dict) -> tuple[bool, str]:
        current = self.value(dash)
        if self.minimum is not None and current < self.minimum:
            return False, f"{self.name}: {self.metric}={current} below {self.minimum}"
        if self.maximum is not None and current > self.maximum:
            return False, f"{self.name}: {self.metric}={current} above {self.maximum}"
        return True, f"{self.name}: ok ({self.metric}={current})"


DEFAULT_ALERTS = [
    AlertRule("latency", "latency_ms.p95", maximum=4000),
    AlertRule("quality: retrieval", "retrieval_hit_rate", minimum=0.80),
    AlertRule("quality: abstention", "abstention_rate", maximum=0.25),
    AlertRule("errors", "error_rate", maximum=0.02),
    AlertRule("degradation", "degraded_rate", maximum=0.10),
    AlertRule("spend", "cost_per_answered_question", maximum=0.05),
]


def evaluate_alerts(records: list[dict], rules: list[AlertRule] = DEFAULT_ALERTS) -> list[str]:
    firing = []
    for rule in rules:
        dash = dashboard(records[-rule.window:])
        ok, message = rule.check(dash)
        if not ok:
            firing.append(message)
    return firing


def compare_windows(records: list[dict], metric: str = "retrieval_hit_rate", window: int = 50) -> dict:
    """Catch a slow regression: this window against the one before it."""
    if len(records) < window * 2:
        return {"enough_data": False}
    recent = dashboard(records[-window:])[metric]
    previous = dashboard(records[-2 * window:-window])[metric]
    return {"enough_data": True, "metric": metric, "previous": previous, "recent": recent,
            "change": round(recent - previous, 3), "regressed": recent < previous - 0.05}


# ---------------------------------------------------------------- replay
def replay(record: dict, run_case) -> dict:
    """Re-run a production failure locally, from its trace, and compare with what happened."""
    output = run_case(record["query"])
    original_ids = record.get("retrieved_ids", [])
    new_ids = output.get("final_ids", [])
    return {"correlation_id": record["correlation_id"], "query": record["query"],
            "original_ids": original_ids, "replay_ids": new_ids,
            "same_retrieval": original_ids == new_ids,
            "original_answer": record.get("answer", "")[:120],
            "replay_answer": str(output.get("answer", ""))[:120]}