"""Metric tracking utilities for SAGE RAG pipelines.

Provides structured metric logging for RAG operations including
retrieval quality, generation metrics, and validation scores.
"""

import logging
import time
from dataclasses import dataclass, field
from typing import Any

from sage.tracking.experiment import get_tracker

logger = logging.getLogger(__name__)


@dataclass
class RAGMetrics:
    """Metrics for a RAG operation.

    Captures comprehensive metrics for evaluating RAG pipeline
    performance including retrieval, generation, and validation.
    """

    # Query metrics
    query_length: int = 0
    query_tokens: int | None = None

    # Retrieval metrics
    retrieval_latency_ms: float = 0.0
    documents_retrieved: int = 0
    documents_used: int = 0
    retrieval_scores: list[float] = field(default_factory=list)
    avg_retrieval_score: float = 0.0

    # Generation metrics
    generation_latency_ms: float = 0.0
    input_tokens: int = 0
    output_tokens: int = 0
    response_length: int = 0

    # Validation metrics
    grounding_score: float | None = None
    citation_accuracy: float | None = None
    safety_passed: bool | None = None

    # Overall metrics
    total_latency_ms: float = 0.0
    success: bool = True
    error: str | None = None

    def to_dict(self) -> dict[str, Any]:
        """Convert metrics to dictionary for logging."""
        metrics = {
            "query_length": self.query_length,
            "retrieval_latency_ms": self.retrieval_latency_ms,
            "documents_retrieved": self.documents_retrieved,
            "documents_used": self.documents_used,
            "avg_retrieval_score": self.avg_retrieval_score,
            "generation_latency_ms": self.generation_latency_ms,
            "input_tokens": self.input_tokens,
            "output_tokens": self.output_tokens,
            "total_tokens": self.input_tokens + self.output_tokens,
            "response_length": self.response_length,
            "total_latency_ms": self.total_latency_ms,
            "success": 1.0 if self.success else 0.0,
        }

        if self.query_tokens is not None:
            metrics["query_tokens"] = self.query_tokens
        if self.grounding_score is not None:
            metrics["grounding_score"] = self.grounding_score
        if self.citation_accuracy is not None:
            metrics["citation_accuracy"] = self.citation_accuracy
        if self.safety_passed is not None:
            metrics["safety_passed"] = 1.0 if self.safety_passed else 0.0

        return metrics

    def calculate_averages(self) -> None:
        """Calculate derived metrics."""
        if self.retrieval_scores:
            self.avg_retrieval_score = sum(self.retrieval_scores) / len(
                self.retrieval_scores
            )


class MetricLogger:
    """Logger for RAG metrics with MLflow integration.

    Provides a convenient interface for logging metrics during
    RAG operations with automatic aggregation and tracking.

    Usage:
        logger = MetricLogger()

        with logger.track_operation("rag-query") as metrics:
            metrics.query_length = len(query)

            with logger.track_phase("retrieval"):
                results = retriever.search(query)
                metrics.documents_retrieved = len(results)

            with logger.track_phase("generation"):
                response = generator.generate(query, results)
                metrics.response_length = len(response)
    """

    def __init__(self, experiment_name: str | None = None):
        """Initialize metric logger.

        Args:
            experiment_name: Optional experiment name for MLflow.
        """
        self._tracker = get_tracker(experiment_name=experiment_name)
        self._current_metrics: RAGMetrics | None = None
        self._phase_start: float | None = None

    def track_operation(
        self,
        operation_name: str,
        params: dict[str, Any] | None = None,
    ) -> "OperationContext":
        """Track a RAG operation.

        Args:
            operation_name: Name of the operation.
            params: Parameters to log.

        Returns:
            OperationContext: Context manager for the operation.
        """
        return OperationContext(self, operation_name, params)

    def track_phase(self, phase_name: str) -> "PhaseContext":
        """Track a phase within an operation.

        Args:
            phase_name: Name of the phase (retrieval, generation, etc.).

        Returns:
            PhaseContext: Context manager for the phase.
        """
        return PhaseContext(self, phase_name)

    def log_metrics(
        self,
        metrics: RAGMetrics,
        run_name: str | None = None,
        params: dict[str, Any] | None = None,
    ) -> None:
        """Log RAG metrics to MLflow.

        Args:
            metrics: RAG metrics to log.
            run_name: Optional run name.
            params: Optional parameters to log.
        """
        metrics.calculate_averages()

        with self._tracker.start_run(run_name=run_name) as run:
            if params:
                run.log_params(params)
            run.log_metrics(metrics.to_dict())


class OperationContext:
    """Context manager for tracking a RAG operation."""

    def __init__(
        self,
        logger: MetricLogger,
        operation_name: str,
        params: dict[str, Any] | None = None,
    ):
        """Initialize operation context.

        Args:
            logger: Parent metric logger.
            operation_name: Name of the operation.
            params: Parameters to log.
        """
        self._logger = logger
        self._operation_name = operation_name
        self._params = params
        self._metrics = RAGMetrics()
        self._start_time: float = 0.0

    def __enter__(self) -> RAGMetrics:
        """Start tracking the operation."""
        self._start_time = time.time()
        self._logger._current_metrics = self._metrics
        return self._metrics

    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
        """End tracking and log metrics."""
        self._metrics.total_latency_ms = (time.time() - self._start_time) * 1000

        if exc_type is not None:
            self._metrics.success = False
            self._metrics.error = str(exc_val)

        self._metrics.calculate_averages()
        self._logger.log_metrics(
            self._metrics,
            run_name=self._operation_name,
            params=self._params,
        )
        self._logger._current_metrics = None


class PhaseContext:
    """Context manager for tracking a phase within an operation."""

    def __init__(self, logger: MetricLogger, phase_name: str):
        """Initialize phase context.

        Args:
            logger: Parent metric logger.
            phase_name: Name of the phase.
        """
        self._logger = logger
        self._phase_name = phase_name
        self._start_time: float = 0.0

    def __enter__(self) -> "PhaseContext":
        """Start tracking the phase."""
        self._start_time = time.time()
        return self

    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
        """End tracking and record latency."""
        latency_ms = (time.time() - self._start_time) * 1000
        metrics = self._logger._current_metrics

        if metrics is not None:
            if self._phase_name == "retrieval":
                metrics.retrieval_latency_ms = latency_ms
            elif self._phase_name == "generation":
                metrics.generation_latency_ms = latency_ms

    @property
    def elapsed_ms(self) -> float:
        """Get elapsed time in milliseconds."""
        return (time.time() - self._start_time) * 1000


def log_rag_metrics(
    query: str,
    response: str,
    sources: list[str],
    retrieval_scores: list[float] | None = None,
    input_tokens: int = 0,
    output_tokens: int = 0,
    retrieval_latency_ms: float = 0.0,
    generation_latency_ms: float = 0.0,
    grounding_score: float | None = None,
    citation_accuracy: float | None = None,
    model: str | None = None,
    run_name: str = "rag-query",
) -> None:
    """Convenience function for logging RAG metrics.

    Args:
        query: User query.
        response: Generated response.
        sources: List of source document IDs.
        retrieval_scores: Retrieval similarity scores.
        input_tokens: Input token count.
        output_tokens: Output token count.
        retrieval_latency_ms: Retrieval latency.
        generation_latency_ms: Generation latency.
        grounding_score: Response grounding score.
        citation_accuracy: Citation accuracy score.
        model: Model identifier.
        run_name: MLflow run name.
    """
    metrics = RAGMetrics(
        query_length=len(query),
        retrieval_latency_ms=retrieval_latency_ms,
        documents_retrieved=len(sources),
        documents_used=len(sources),
        retrieval_scores=retrieval_scores or [],
        generation_latency_ms=generation_latency_ms,
        input_tokens=input_tokens,
        output_tokens=output_tokens,
        response_length=len(response),
        total_latency_ms=retrieval_latency_ms + generation_latency_ms,
        grounding_score=grounding_score,
        citation_accuracy=citation_accuracy,
    )

    params = {"model": model} if model else {}

    logger = MetricLogger()
    logger.log_metrics(metrics, run_name=run_name, params=params)
