from __future__ import annotations

from collections import defaultdict, deque
from dataclasses import dataclass, field


@dataclass
class Stats:
    successes: int = 0
    failures: int = 0
    latencies_ms: deque[float] = field(default_factory=lambda: deque(maxlen=512))

    def record(self, ok: bool, latency_ms: float | None = None) -> None:
        if ok:
            self.successes += 1
        else:
            self.failures += 1
        if latency_ms is not None:
            self.latencies_ms.append(latency_ms)

    def p50(self) -> float:
        vals = sorted(self.latencies_ms)
        if not vals:
            return 0.0
        return vals[len(vals) // 2]

    def p95(self) -> float:
        vals = sorted(self.latencies_ms)
        if not vals:
            return 0.0
        idx = int(len(vals) * 0.95) - 1
        idx = max(0, min(idx, len(vals) - 1))
        return vals[idx]


class MetricsRegistry:
    def __init__(self) -> None:
        self._by_provider: dict[str, Stats] = {}
        self._auth: dict[str, dict[str, int]] = defaultdict(lambda: {"success": 0, "failure": 0})
        self._auth_fail_reasons: dict[str, int] = defaultdict(int)
        self._billing_events: dict[str, int] = defaultdict(int)
        self._billing_finalize_lag: deque[float] = deque(maxlen=512)
        self._usage_totals: dict[str, dict[str, int]] = defaultdict(
            lambda: {"prompt": 0, "completion": 0}
        )
        # Adaptive concurrency metrics
        self._concurrency_adjustments: dict[str, dict[str, int]] = defaultdict(
            lambda: {"increases": 0, "decreases": 0, "rate_limits": 0}
        )
        self._concurrency_limits: dict[str, float] = {}

    def for_provider(self, name: str) -> Stats:
        st = self._by_provider.get(name)
        if not st:
            st = Stats()
            self._by_provider[name] = st
        return st

    def snapshot(self) -> dict[str, dict[str, float]]:
        snap: dict[str, dict[str, float]] = {}
        for k, st in self._by_provider.items():
            total = st.successes + st.failures
            success_rate = (st.successes / total * 100.0) if total > 0 else 0.0
            snap[k] = {
                "successes": float(st.successes),
                "failures": float(st.failures),
                "success_rate_pct": success_rate,
                "p50_ms": st.p50(),
                "p95_ms": st.p95(),
            }
        snap["auth"] = {
            backend: float(counts["success"] + counts["failure"])
            for backend, counts in self._auth.items()
        }
        snap["billing_events"] = {k: float(v) for k, v in self._billing_events.items()}
        snap["billing_finalize"] = {
            "samples": float(len(self._billing_finalize_lag)),
            "p50_s": _percentile(self._billing_finalize_lag, 0.5),
            "p95_s": _percentile(self._billing_finalize_lag, 0.95),
        }
        return snap

    def render_text(self) -> str:
        """Render a minimal text metrics format.

        This is intentionally generic and not tightly coupled to Prometheus.
        """
        lines = []
        for k, st in self._by_provider.items():
            total = st.successes + st.failures
            success_rate = (st.successes / total * 100.0) if total > 0 else 0.0
            lines.append(f'takumi_successes{{provider="{k}"}} {st.successes}')
            lines.append(f'takumi_failures{{provider="{k}"}} {st.failures}')
            lines.append(f'takumi_p50_ms{{provider="{k}"}} {st.p50():.3f}')
            lines.append(f'takumi_p95_ms{{provider="{k}"}} {st.p95():.3f}')
            lines.append(f'takumi_success_rate_pct{{provider="{k}"}} {success_rate:.3f}')
        for backend, counts in self._auth.items():
            lines.append(f'takumi_auth_success_total{{backend="{backend}"}} {counts["success"]}')
            lines.append(f'takumi_auth_failure_total{{backend="{backend}"}} {counts["failure"]}')
        for key, value in self._auth_fail_reasons.items():
            backend, reason = key.split(":", 1)
            lines.append(
                f'takumi_auth_failure_reason_total{{backend="{backend}",reason="{reason}"}} {value}'
            )
        for result, value in self._billing_events.items():
            lines.append(f'takumi_billing_events_total{{result="{result}"}} {value}')
        if self._billing_finalize_lag:
            lines.append(
                "takumi_billing_finalize_lag_p50_seconds "
                f"{_percentile(self._billing_finalize_lag, 0.5):.3f}"
            )
            lines.append(
                "takumi_billing_finalize_lag_p95_seconds "
                f"{_percentile(self._billing_finalize_lag, 0.95):.3f}"
            )
        for key, totals in self._usage_totals.items():
            org, model = key.split(":", 1)
            lines.append(
                "takumi_usage_prompt_tokens_total"
                f'{{org_id="{org}",model="{model}"}} {totals["prompt"]}'
            )
            lines.append(
                "takumi_usage_completion_tokens_total"
                f'{{org_id="{org}",model="{model}"}} {totals["completion"]}'
            )
        # Adaptive concurrency metrics
        for provider, adjustments in self._concurrency_adjustments.items():
            lines.append(
                f'takumi_concurrency_increases_total{{provider="{provider}"}} '
                f'{adjustments["increases"]}'
            )
            lines.append(
                f'takumi_concurrency_decreases_total{{provider="{provider}"}} '
                f'{adjustments["decreases"]}'
            )
            lines.append(
                f'takumi_concurrency_rate_limits_total{{provider="{provider}"}} '
                f'{adjustments["rate_limits"]}'
            )
        for provider, limit in self._concurrency_limits.items():
            lines.append(
                f'takumi_concurrency_limit{{provider="{provider}"}} {limit:.1f}'
            )
        return "\n".join(lines) + "\n"

    def record_auth(self, backend: str, ok: bool, *, reason: str | None = None) -> None:
        entry = self._auth[backend]
        if ok:
            entry["success"] += 1
        else:
            entry["failure"] += 1
            if reason:
                key = f"{backend}:{reason}"
                self._auth_fail_reasons[key] += 1

    def record_billing_event(self, result: str) -> None:
        self._billing_events[result] += 1

    def record_finalize_lag(self, seconds: float) -> None:
        if seconds >= 0:
            self._billing_finalize_lag.append(seconds)

    def record_usage_tokens(self, org_id: str, model: str, prompt: int, completion: int) -> None:
        key = f"{org_id}:{model}"
        entry = self._usage_totals[key]
        entry["prompt"] += max(0, int(prompt))
        entry["completion"] += max(0, int(completion))

    def record_concurrency_adjustment(
        self, provider: str, old_limit: float, new_limit: float, is_rate_limit: bool
    ) -> None:
        """Record a concurrency limit adjustment from the adaptive limiter.

        Args:
            provider: The provider key (e.g., "openai", "anthropic").
            old_limit: Previous concurrency limit.
            new_limit: New concurrency limit after adjustment.
            is_rate_limit: True if this was triggered by a 429 response.
        """
        entry = self._concurrency_adjustments[provider]
        if new_limit > old_limit:
            entry["increases"] += 1
        elif new_limit < old_limit:
            entry["decreases"] += 1
        if is_rate_limit:
            entry["rate_limits"] += 1
        self._concurrency_limits[provider] = new_limit

    def get_concurrency_limits(self) -> dict[str, float]:
        """Get current concurrency limits for all tracked providers."""
        return dict(self._concurrency_limits)


def _percentile(values: deque[float], frac: float) -> float:
    if not values:
        return 0.0
    ordered = sorted(values)
    idx = int(len(ordered) * frac)
    idx = max(0, min(idx, len(ordered) - 1))
    return ordered[idx]


METRICS = MetricsRegistry()
