"""Telemetry, billing, and usage-tracking helpers for gateway HTTP handlers."""

from __future__ import annotations

import json
import logging
from typing import Any, Mapping

from ember.models.schemas import ChatResponse as ProviderChatResponse
from ember.models.schemas import UsageStats

from ..domain.model import ChatResponse, StreamingResponse, Usage
from ..infra.metrics import METRICS as PROVIDER_METRICS
from ..infra.sur import generate_sur
from ..pipeline.container import Services
from ..pipeline.pipeline import Context

LOG = logging.getLogger(__name__)
STREAM_LOG = logging.getLogger("ember.gateway.stream")

# EB1 synthesis fields that should only be visible to admin/operator/internal roles.
# These are stripped from responses_payload.metadata for regular users.
_SYNTHESIS_DEBUG_FIELDS = frozenset(
    {
        "synthesis_model",
        "primary_source_model",
        "synthesized_response",
        "synthesis_rationale",
        "synthesis_approach",
    }
)


def _emit_request_telemetry(ctx: Context, response: ChatResponse) -> None:
    """Emit a single structured telemetry event and feed provider-level metrics.

    This is the SSOT log event for per-request latency breakdown. It captures
    pipeline step timings, provider routing decisions, usage, and cost in one
    structured JSON line that downstream tooling (jq, pandas, dashboards) can
    consume directly.
    """
    finish_reason = (
        response.choices[0].finish_reason if response.choices else None
    )
    LOG.info(
        "request_telemetry",
        extra={
            "request_id": ctx.request_id,
            "org_id": ctx.req.org_id,
            "model": ctx.canonical_model,
            "provider": ctx.chosen_provider,
            "provider_model": ctx.provider_model,
            "stream": ctx.req.stream,
            "prompt_tokens": response.usage.prompt_tokens,
            "completion_tokens": response.usage.completion_tokens,
            "cost_usd": response.usage.cost_usd,
            "finish_reason": finish_reason,
            "timings_ms": ctx.step_timings_ms,
            "pipeline_total_ms": ctx.step_timings_ms.get("total_ms", 0),
        },
    )
    if ctx.chosen_provider:
        PROVIDER_METRICS.for_provider(ctx.chosen_provider).record(
            ok=True,
            latency_ms=ctx.step_timings_ms.get("step_execute", 0),
        )


def _persist_sur(
    services: Services,
    *,
    request_id: str,
    org_id: str,
    price_book_version: str,
    rbc_id: str | None,
    provider_prompt_tokens: int,
    provider_completion_tokens: int,
    canonical_prompt_tokens: int,
    canonical_completion_tokens: int,
) -> None:
    sur = generate_sur(
        request_id=request_id,
        org_id=org_id,
        price_book_version=price_book_version,
        rbc_id=rbc_id,
        provider_prompt_tokens=provider_prompt_tokens,
        provider_completion_tokens=provider_completion_tokens,
        canonical_prompt_tokens=canonical_prompt_tokens,
        canonical_completion_tokens=canonical_completion_tokens,
    )
    services.sur.persist(sur)


def _finalize_stream_usage(
    services: Services,
    response: StreamingResponse,
    *,
    recorded_prompt_tokens: int,
    recorded_completion_tokens: int,
    recorded_cost: float,
) -> Usage | None:
    stream_source = response.stream_source
    final_payload = getattr(stream_source, "final_response", None) if stream_source else None

    prompt_tokens, completion_tokens, total_tokens = _extract_usage_totals(
        final_payload,
        fallback_prompt=response.prompt_estimate,
    )

    remaining_prompt = max(0, prompt_tokens - recorded_prompt_tokens)
    remaining_completion = max(0, completion_tokens - recorded_completion_tokens)

    # Sanity check: recorded tokens should never exceed final reported tokens.
    # If this fires, there's a bug in chunk-level usage accumulation.
    if recorded_completion_tokens > completion_tokens:
        STREAM_LOG.warning(
            "Token accounting mismatch: recorded %d completion tokens > final %d",
            recorded_completion_tokens,
            completion_tokens,
        )

    if remaining_completion and not services.tpm.consume_output(
        response.org_id, remaining_completion
    ):
        raise services.errors.rate_limited(scope="tpm", phase="completion_postcharge")

    total_tokens = total_tokens or (prompt_tokens + completion_tokens)
    usage = Usage(
        prompt_tokens=prompt_tokens,
        completion_tokens=completion_tokens,
        total_tokens=total_tokens,
        cost_usd=0.0,
    )

    total_cost = float(
        services.pricing.calculate(
            model=response.provider_model,
            prompt_tokens=usage.prompt_tokens,
            completion_tokens=usage.completion_tokens,
            price_book_version=response.price_book_version,
            billing_mode=response.billing_mode,
        )
    )
    usage.cost_usd = total_cost

    remaining_cost = max(0.0, total_cost - recorded_cost)

    services.billing.record_delta(
        response.fbr_id,
        remaining_prompt,
        remaining_completion,
        remaining_cost,
    )
    services.billing.finalize(response.fbr_id)
    _persist_sur(
        services,
        request_id=response.id,
        org_id=response.org_id,
        price_book_version=response.price_book_version,
        rbc_id=getattr(response.rbc, "id", None),
        provider_prompt_tokens=usage.prompt_tokens,
        provider_completion_tokens=usage.completion_tokens,
        canonical_prompt_tokens=usage.prompt_tokens,
        canonical_completion_tokens=usage.completion_tokens,
    )
    return usage


def _extract_usage_totals(
    final_payload: Any,
    *,
    fallback_prompt: int,
) -> tuple[int, int, int]:
    prompt_tokens = int(max(0, fallback_prompt))
    completion_tokens = 0
    total_tokens = prompt_tokens

    if isinstance(final_payload, ProviderChatResponse):
        usage = final_payload.usage or UsageStats()
        prompt_tokens = int(getattr(usage, "prompt_tokens", prompt_tokens) or prompt_tokens)
        completion_tokens = int(getattr(usage, "completion_tokens", 0) or 0)
        total_tokens = int(
            getattr(usage, "total_tokens", prompt_tokens + completion_tokens)
            or (prompt_tokens + completion_tokens)
        )
    elif isinstance(final_payload, Mapping):
        usage = final_payload.get("usage")
        if isinstance(usage, Mapping):
            prompt_tokens = int(usage.get("prompt_tokens") or prompt_tokens)
            completion_tokens = int(usage.get("completion_tokens") or 0)
            total_tokens = int(usage.get("total_tokens") or (prompt_tokens + completion_tokens))

    total_tokens = max(total_tokens, prompt_tokens + completion_tokens)
    return prompt_tokens, completion_tokens, total_tokens


def _sse_event(event: str, data: Any, event_id: int | None = None) -> bytes:
    payload = json.dumps(data)
    prefix = ""
    if event_id is not None:
        prefix = f"id: {event_id}\n"
    return f"{prefix}event: {event}\ndata: {payload}\n\n".encode("utf-8")


def _extract_usage_delta(data: str) -> tuple[int, int] | None:
    try:
        parsed = json.loads(data)
    except json.JSONDecodeError:
        return None

    if not isinstance(parsed, Mapping):
        return None

    candidates = [parsed.get("usage_delta"), parsed.get("usage"), parsed.get("delta")]
    for candidate in candidates:
        if isinstance(candidate, Mapping):
            prompt = int(candidate.get("prompt_tokens") or candidate.get("prompt") or 0)
            completion = int(candidate.get("completion_tokens") or candidate.get("completion") or 0)
            if prompt or completion:
                return prompt, completion
    return None


def _extract_rbc_id(raw: Any) -> str | None:
    if isinstance(raw, Mapping):
        candidate = raw.get("rbc_id") or raw.get("id")
        if candidate is None:
            # Legacy embedding sometimes nests under `rbc` key.
            nested = raw.get("rbc")
            if isinstance(nested, Mapping):
                candidate = nested.get("id")
        if candidate is not None:
            return str(candidate)
    return None
