"""Concise model connection and latency diagnostics for production logs."""

from __future__ import annotations

import logging
import threading
import time
from dataclasses import dataclass
from typing import Any
from uuid import UUID

from langchain_core.callbacks import BaseCallbackHandler

logger = logging.getLogger(__name__)


@dataclass(slots=True)
class _CallTiming:
    started: float
    first_chunk: float | None = None


class ModelTimingCallback(BaseCallbackHandler):
    """Log model-call start, first response chunk, completion, and failure."""

    def __init__(self, *, model_name: str, base_url: str | None) -> None:
        self._model_name = model_name
        self._base_url = base_url or "-"
        self._calls: dict[str, _CallTiming] = {}
        self._lock = threading.Lock()

    def on_chat_model_start(
        self,
        serialized: dict[str, Any],
        messages: list[list[Any]],
        *,
        run_id: UUID,
        **kwargs: Any,
    ) -> None:
        rid = str(run_id)
        now = time.monotonic()
        with self._lock:
            self._calls[rid] = _CallTiming(started=now)
        logger.info(
            "LLM_CONNECT_START model=%s base_url=%s run_id=%s message_batches=%d message_count=%d",
            self._model_name,
            self._base_url,
            rid,
            len(messages),
            sum(len(batch) for batch in messages),
        )

    def on_llm_new_token(
        self,
        token: str,
        *,
        run_id: UUID,
        chunk: Any | None = None,
        **kwargs: Any,
    ) -> None:
        """Treat the first callback as connection success, including reasoning-only chunks."""
        rid = str(run_id)
        now = time.monotonic()
        with self._lock:
            timing = self._calls.get(rid)
            if timing is None or timing.first_chunk is not None:
                return
            timing.first_chunk = now
            first_chunk_ms = int((now - timing.started) * 1000)
        logger.info(
            "LLM_CONNECT_OK model=%s base_url=%s run_id=%s first_chunk_ms=%d",
            self._model_name,
            self._base_url,
            rid,
            first_chunk_ms,
        )

    def on_llm_end(self, response: Any, *, run_id: UUID, **kwargs: Any) -> None:
        self._finish(run_id=run_id, error=None)

    def on_llm_error(self, error: BaseException, *, run_id: UUID, **kwargs: Any) -> None:
        self._finish(run_id=run_id, error=error)

    def _finish(self, *, run_id: UUID, error: BaseException | None) -> None:
        rid = str(run_id)
        now = time.monotonic()
        with self._lock:
            timing = self._calls.pop(rid, None)

        if timing is None:
            logger.error(
                "LLM_REQUEST_FAILED model=%s base_url=%s run_id=%s error_type=%s error=%s",
                self._model_name,
                self._base_url,
                rid,
                type(error).__name__ if error else "UnknownLifecycle",
                error or "missing start event",
            )
            return

        total_ms = int((now - timing.started) * 1000)
        first_chunk_ms = int((timing.first_chunk - timing.started) * 1000) if timing.first_chunk else None
        if error is None:
            logger.info(
                "LLM_REQUEST_DONE model=%s base_url=%s run_id=%s status=success first_chunk_ms=%s total_ms=%d",
                self._model_name,
                self._base_url,
                rid,
                first_chunk_ms if first_chunk_ms is not None else "none",
                total_ms,
            )
            return

        logger.error(
            "LLM_REQUEST_FAILED model=%s base_url=%s run_id=%s status=error first_chunk_ms=%s total_ms=%d error_type=%s error=%s",
            self._model_name,
            self._base_url,
            rid,
            first_chunk_ms if first_chunk_ms is not None else "none",
            total_ms,
            type(error).__name__,
            error,
        )