"""Trace service for LLM observability.

Coordinates the recording, querying, and analysis of ``LLMTrace`` entries.
Optionally forwards traces to LangSmith when the ``LANGCHAIN_TRACING_V2``
environment variable is set to ``"true"``.

Metric collection hooks produce ``MetricEntry`` objects from aggregated
trace data for a given plan.

Design decisions:
    - Dependency injection: ``Settings`` and ``LLMTraceRepository`` are
      injected via the constructor (consistent with existing service patterns).
    - Stateless: All trace state lives in the repository.
    - LangSmith forwarding is best-effort — failures are logged but do not
      raise.

Based on Forgejo issue #500.
"""

from __future__ import annotations

import logging
import os
from typing import TYPE_CHECKING, Any

from cleveragents.domain.models.observability.llm_trace import LLMTrace
from cleveragents.domain.models.observability.metrics import (
    MetricCollector,
    MetricEntry,
    OperationalMetricKey,
)

if TYPE_CHECKING:
    from cleveragents.config.settings import Settings
    from cleveragents.infrastructure.database.llm_trace_repository import (
        LLMTraceRepository,
    )

logger = logging.getLogger(__name__)


class TraceService:
    """Orchestrates LLM trace persistence, querying, and metric computation.

    Attributes:
        _settings: Application settings (injected).
        _repository: Persistence layer for ``LLMTrace`` records.
    """

    def __init__(
        self,
        settings: Settings,
        repository: LLMTraceRepository,
    ) -> None:
        """Initialise with injected dependencies.

        Args:
            settings: Application settings.
            repository: LLM trace repository for persistence.
        """
        self._settings = settings
        self._repository = repository

    # ---- write path -------------------------------------------------------

    def record_trace(self, trace: LLMTrace) -> None:
        """Persist an ``LLMTrace`` entry and optionally forward to LangSmith.

        Args:
            trace: The trace to record.
        """
        self._repository.save(trace)
        logger.debug("Recorded LLM trace %s for plan %s", trace.trace_id, trace.plan_id)

        # Best-effort LangSmith forwarding
        if self._langsmith_enabled():
            self.forward_to_langsmith(trace)

    # ---- read path --------------------------------------------------------

    def get_traces(self, plan_id: str) -> list[LLMTrace]:
        """Query all traces for a plan.

        Args:
            plan_id: ULID of the plan.

        Returns:
            List of ``LLMTrace`` objects ordered by timestamp.
        """
        return self._repository.list_by_plan(plan_id)

    def get_traces_by_decision(self, decision_id: str) -> list[LLMTrace]:
        """Query all traces for a decision.

        Args:
            decision_id: ULID of the decision.

        Returns:
            List of ``LLMTrace`` objects ordered by timestamp.
        """
        return self._repository.list_by_decision(decision_id)

    def get_trace(self, trace_id: str) -> LLMTrace | None:
        """Retrieve a single trace by ID.

        Args:
            trace_id: ULID of the trace.

        Returns:
            The ``LLMTrace`` or ``None``.
        """
        return self._repository.get(trace_id)

    # ---- metric computation -----------------------------------------------

    def compute_metrics(self, plan_id: str) -> list[MetricEntry]:
        """Compute operational metrics from all traces for a plan.

        Aggregates LLM-level metrics (call count, total tokens, total
        cost, average latency) from the trace entries.

        Args:
            plan_id: ULID of the plan.

        Returns:
            List of ``MetricEntry`` objects for the computed metrics.
        """
        traces = self._repository.list_by_plan(plan_id)
        metrics: list[MetricEntry] = []

        if not traces:
            return metrics

        total_tokens = sum(t.prompt_tokens + t.completion_tokens for t in traces)
        total_cost = sum(t.cost_usd for t in traces)
        total_latency = sum(t.latency_ms for t in traces)
        call_count = len(traces)
        avg_latency = total_latency / call_count if call_count > 0 else 0.0

        metrics.append(
            MetricCollector.record(
                OperationalMetricKey.LLM_CALL_COUNT,
                float(call_count),
                plan_id,
            )
        )
        metrics.append(
            MetricCollector.record(
                OperationalMetricKey.LLM_TOTAL_TOKENS,
                float(total_tokens),
                plan_id,
            )
        )
        metrics.append(
            MetricCollector.record(
                OperationalMetricKey.LLM_TOTAL_COST_USD,
                total_cost,
                plan_id,
            )
        )
        metrics.append(
            MetricCollector.record(
                OperationalMetricKey.LLM_AVG_LATENCY_MS,
                avg_latency,
                plan_id,
            )
        )

        return metrics

    # ---- plan lifecycle hooks (stubs) -------------------------------------

    def on_plan_start(self, plan_id: str) -> MetricEntry | None:
        """Hook: called when a plan begins execution.

        Returns a stub ``MetricEntry`` for plan decision count
        initialisation.  Real timing starts here.

        Args:
            plan_id: ULID of the plan.

        Returns:
            A ``MetricEntry`` for ``PLAN_DECISION_COUNT`` initialised to 0.
        """
        return MetricCollector.record(
            OperationalMetricKey.PLAN_DECISION_COUNT,
            0.0,
            plan_id,
        )

    def on_actor_invocation(
        self, plan_id: str, actor: str, latency_ms: float
    ) -> list[MetricEntry]:
        """Hook: called after an actor invocation completes.

        Args:
            plan_id: ULID of the plan.
            actor: Name of the invoked actor.
            latency_ms: Actor invocation latency in milliseconds.

        Returns:
            List of ``MetricEntry`` objects for actor metrics.
        """
        return [
            MetricCollector.record(
                OperationalMetricKey.ACTOR_INVOCATION_COUNT,
                1.0,
                plan_id,
                labels={"actor": actor},
            ),
            MetricCollector.record(
                OperationalMetricKey.ACTOR_LATENCY_MS,
                latency_ms,
                plan_id,
                labels={"actor": actor},
            ),
        ]

    def on_tool_execution(
        self, plan_id: str, tool_name: str, *, errored: bool = False
    ) -> list[MetricEntry]:
        """Hook: called after a tool execution completes.

        Args:
            plan_id: ULID of the plan.
            tool_name: Name of the tool executed.
            errored: Whether the execution errored.

        Returns:
            List of ``MetricEntry`` objects for tool metrics.
        """
        entries = [
            MetricCollector.record(
                OperationalMetricKey.TOOL_INVOCATION_COUNT,
                1.0,
                plan_id,
                labels={"tool": tool_name},
            ),
        ]
        if errored:
            entries.append(
                MetricCollector.record(
                    OperationalMetricKey.TOOL_ERROR_RATE,
                    1.0,
                    plan_id,
                    labels={"tool": tool_name},
                )
            )
        return entries

    # ---- LangSmith forwarding ---------------------------------------------

    @staticmethod
    def _langsmith_enabled() -> bool:
        """Check whether LangSmith forwarding is enabled."""
        return os.environ.get("LANGCHAIN_TRACING_V2", "").lower() == "true"

    def forward_to_langsmith(self, trace: LLMTrace) -> None:
        """Forward a trace to LangSmith (best-effort).

        This is a best-effort operation: failures are logged at
        warning level but do not propagate exceptions.

        Args:
            trace: The trace to forward.
        """
        try:
            _forward_trace_to_langsmith(trace)
        except Exception:
            logger.warning(
                "Failed to forward trace %s to LangSmith",
                trace.trace_id,
                exc_info=True,
            )


def _forward_trace_to_langsmith(trace: LLMTrace) -> None:
    """Internal: forward a trace to LangSmith via the langsmith SDK.

    Separated for testability.  Imports ``langsmith`` lazily to avoid
    hard dependency.

    Args:
        trace: The trace to forward.
    """
    try:
        from langsmith import Client  # type: ignore[import-untyped]
    except ImportError:
        logger.debug("langsmith SDK not installed; skipping trace forwarding")
        return

    client = Client()
    run_data: dict[str, Any] = {
        "name": f"llm-trace-{trace.trace_id}",
        "run_type": "llm",
        "inputs": {
            "actor": trace.actor,
            "provider": trace.provider,
            "model": trace.model,
            "prompt_tokens": trace.prompt_tokens,
            "streaming": trace.streaming,
        },
        "outputs": {
            "completion_tokens": trace.completion_tokens,
            "cost_usd": trace.cost_usd,
            "latency_ms": trace.latency_ms,
            "tool_calls": trace.tool_calls,
        },
    }
    if trace.error:
        run_data["error"] = trace.error

    client.create_run(**run_data)
    logger.debug("Forwarded trace %s to LangSmith", trace.trace_id)
