"""LLM trace persistence repository.

Provides ``LLMTraceRepository`` for CRUD operations on the
``llm_traces`` table.  Follows the session-factory pattern (ADR-007)
with retry logic (ADR-033).

Based on Forgejo issue #500.
"""

from __future__ import annotations

import json
from collections.abc import Callable
from datetime import datetime
from typing import Any

from sqlalchemy.exc import DatabaseError as SQLAlchemyDatabaseError
from sqlalchemy.exc import OperationalError
from sqlalchemy.orm import Session

from cleveragents.core.exceptions import DatabaseError
from cleveragents.core.retry_patterns import (
    retry_database_operation as database_retry,
)
from cleveragents.domain.models.observability.llm_trace import LLMTrace
from cleveragents.infrastructure.database.models import LLMTraceModel


class LLMTraceRepository:
    """Repository for persisting and querying ``LLMTrace`` records.

    Uses the session-factory pattern: each public method obtains a
    session from the factory.  Callers are responsible for commit.
    """

    def __init__(
        self,
        session_factory: Callable[[], Session],
    ) -> None:
        """Initialise with a callable returning a SQLAlchemy Session."""
        if session_factory is None:
            raise ValueError("session_factory must not be None")
        self._sf = session_factory

    def _session(self) -> Session:
        return self._sf()

    @database_retry
    def save(self, trace: LLMTrace) -> None:
        """Persist a single ``LLMTrace`` row.

        Args:
            trace: The trace to persist.

        Raises:
            DatabaseError: On unrecoverable persistence failure.
        """
        session = self._session()
        try:
            model = LLMTraceModel(
                trace_id=trace.trace_id,
                plan_id=trace.plan_id,
                decision_id=trace.decision_id,
                actor=trace.actor,
                provider=trace.provider,
                model=trace.model,
                prompt_tokens=trace.prompt_tokens,
                completion_tokens=trace.completion_tokens,
                cost_usd=trace.cost_usd,
                latency_ms=trace.latency_ms,
                tool_calls_json=json.dumps(trace.tool_calls)
                if trace.tool_calls
                else None,
                context_hash=trace.context_hash,
                streaming=trace.streaming,
                retry_count=trace.retry_count,
                error=trace.error,
                timestamp=trace.timestamp.isoformat(),
            )
            session.add(model)
            session.commit()
        except (SQLAlchemyDatabaseError, OperationalError) as exc:
            session.rollback()
            raise DatabaseError(f"Failed to save LLM trace: {exc}") from exc

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

        Args:
            trace_id: ULID of the trace to retrieve.

        Returns:
            The ``LLMTrace`` or ``None`` if not found.
        """
        session = self._session()
        try:
            row = session.query(LLMTraceModel).filter_by(trace_id=trace_id).first()
            if row is None:
                return None
            return self._to_domain(row)
        except (SQLAlchemyDatabaseError, OperationalError) as exc:
            raise DatabaseError(f"Failed to get LLM trace: {exc}") from exc

    @database_retry
    def list_by_plan(self, plan_id: str) -> list[LLMTrace]:
        """List all traces for a given plan.

        Args:
            plan_id: ULID of the plan.

        Returns:
            List of ``LLMTrace`` objects ordered by timestamp.
        """
        session = self._session()
        try:
            rows = (
                session.query(LLMTraceModel)
                .filter_by(plan_id=plan_id)
                .order_by(LLMTraceModel.timestamp)
                .all()
            )
            return [self._to_domain(r) for r in rows]
        except (SQLAlchemyDatabaseError, OperationalError) as exc:
            raise DatabaseError(f"Failed to list traces by plan: {exc}") from exc

    @database_retry
    def list_by_decision(self, decision_id: str) -> list[LLMTrace]:
        """List all traces for a given decision.

        Args:
            decision_id: ULID of the decision.

        Returns:
            List of ``LLMTrace`` objects ordered by timestamp.
        """
        session = self._session()
        try:
            rows = (
                session.query(LLMTraceModel)
                .filter_by(decision_id=decision_id)
                .order_by(LLMTraceModel.timestamp)
                .all()
            )
            return [self._to_domain(r) for r in rows]
        except (SQLAlchemyDatabaseError, OperationalError) as exc:
            raise DatabaseError(f"Failed to list traces by decision: {exc}") from exc

    @staticmethod
    def _to_domain(row: Any) -> LLMTrace:
        """Convert a database row to a domain ``LLMTrace``."""
        tool_calls: list[dict[str, Any]] = []
        if row.tool_calls_json:
            tool_calls = json.loads(row.tool_calls_json)

        return LLMTrace(
            trace_id=row.trace_id,
            plan_id=row.plan_id,
            decision_id=row.decision_id,
            actor=row.actor,
            provider=row.provider,
            model=row.model,
            prompt_tokens=row.prompt_tokens,
            completion_tokens=row.completion_tokens,
            cost_usd=row.cost_usd,
            latency_ms=row.latency_ms,
            tool_calls=tool_calls,
            context_hash=row.context_hash,
            streaming=row.streaming,
            retry_count=row.retry_count,
            error=row.error,
            timestamp=datetime.fromisoformat(row.timestamp),
        )
