"""MLflow Experiment Tracking for SAGE.

Provides utilities for tracking RAG experiments, including
parameters, metrics, and artifacts.
"""

import functools
import logging
import os
import time
from collections.abc import Callable, Generator
from contextlib import contextmanager
from dataclasses import dataclass, field
from datetime import UTC, datetime
from typing import Any, TypeVar

logger = logging.getLogger(__name__)

# Type variable for decorated functions
F = TypeVar("F", bound=Callable[..., Any])

# Check if MLflow is available
try:
    from mlflow.tracking import MlflowClient

    import mlflow

    MLFLOW_AVAILABLE = True
except ImportError:
    mlflow = None  # type: ignore[assignment]
    MlflowClient = None  # type: ignore[assignment, misc]
    MLFLOW_AVAILABLE = False


@dataclass
class ExperimentConfig:
    """Configuration for experiment tracking."""

    experiment_name: str = "sage-experiments"
    tracking_uri: str = "http://localhost:5000"
    artifact_location: str | None = None
    tags: dict[str, str] = field(default_factory=dict)
    enabled: bool = True


class ExperimentTracker:
    """MLflow experiment tracker for SAGE.

    Provides a high-level interface for tracking experiments,
    logging metrics, and managing runs.

    Usage:
        tracker = ExperimentTracker()

        with tracker.start_run("query-evaluation"):
            tracker.log_params({"model": "gpt-4", "temperature": 0.7})
            tracker.log_metrics({"accuracy": 0.95, "latency_ms": 150})
            tracker.log_artifact("results.json")
    """

    def __init__(
        self,
        config: ExperimentConfig | None = None,
        experiment_name: str | None = None,
        tracking_uri: str | None = None,
    ):
        """Initialize experiment tracker.

        Args:
            config: Experiment configuration.
            experiment_name: Override experiment name.
            tracking_uri: Override tracking URI.
        """
        self.config = config or ExperimentConfig()

        if experiment_name:
            self.config.experiment_name = experiment_name
        if tracking_uri:
            self.config.tracking_uri = tracking_uri

        self._client: MlflowClient | None = None
        self._experiment_id: str | None = None
        self._active_run: Any = None

        if MLFLOW_AVAILABLE and self.config.enabled:
            self._initialize_mlflow()
        elif not MLFLOW_AVAILABLE:
            logger.warning("MLflow not installed. Experiment tracking disabled.")
            self.config.enabled = False

    def _initialize_mlflow(self) -> None:
        """Initialize MLflow client and experiment."""
        if not MLFLOW_AVAILABLE:
            return

        try:
            mlflow.set_tracking_uri(self.config.tracking_uri)  # type: ignore[union-attr]
            self._client = MlflowClient()

            # Create or get experiment
            experiment = mlflow.get_experiment_by_name(self.config.experiment_name)  # type: ignore[union-attr]
            if experiment is None:
                self._experiment_id = mlflow.create_experiment(  # type: ignore[union-attr]
                    self.config.experiment_name,
                    artifact_location=self.config.artifact_location,
                    tags=self.config.tags,
                )
            else:
                self._experiment_id = experiment.experiment_id

            mlflow.set_experiment(self.config.experiment_name)  # type: ignore[union-attr]
            logger.info(
                f"MLflow initialized: experiment={self.config.experiment_name}, "
                f"tracking_uri={self.config.tracking_uri}"
            )
        except Exception as e:
            logger.warning(f"Failed to initialize MLflow: {e}")
            self.config.enabled = False

    @property
    def is_enabled(self) -> bool:
        """Check if tracking is enabled."""
        return self.config.enabled and MLFLOW_AVAILABLE

    @contextmanager
    def start_run(
        self,
        run_name: str | None = None,
        tags: dict[str, str] | None = None,
        nested: bool = False,
    ) -> Generator["RunContext", None, None]:
        """Start a new experiment run.

        Args:
            run_name: Name for the run.
            tags: Additional tags for the run.
            nested: Whether this is a nested run.

        Yields:
            RunContext: Context manager for the run.
        """
        if not self.is_enabled:
            yield RunContext(self, None)
            return

        all_tags = {**self.config.tags, **(tags or {})}
        all_tags["sage.version"] = os.environ.get("SAGE_VERSION", "dev")
        all_tags["sage.timestamp"] = datetime.now(UTC).isoformat()

        try:
            self._active_run = mlflow.start_run(  # type: ignore[union-attr]
                run_name=run_name,
                tags=all_tags,
                nested=nested,
            )
            yield RunContext(self, self._active_run)
        finally:
            if self._active_run:
                mlflow.end_run()  # type: ignore[union-attr]
                self._active_run = None

    def log_params(self, params: dict[str, Any]) -> None:
        """Log parameters for the current run.

        Args:
            params: Dictionary of parameters to log.
        """
        if not self.is_enabled or not self._active_run:
            return

        # Convert values to strings and truncate if needed
        safe_params = {}
        for key, value in params.items():
            str_value = str(value)
            # MLflow has a 500 char limit for param values
            if len(str_value) > 500:
                str_value = str_value[:497] + "..."
            safe_params[key] = str_value

        try:
            mlflow.log_params(safe_params)  # type: ignore[union-attr]
        except Exception as e:
            logger.warning(f"Failed to log params: {e}")

    def log_metrics(
        self,
        metrics: dict[str, float],
        step: int | None = None,
    ) -> None:
        """Log metrics for the current run.

        Args:
            metrics: Dictionary of metrics to log.
            step: Optional step number for the metrics.
        """
        if not self.is_enabled or not self._active_run:
            return

        try:
            mlflow.log_metrics(metrics, step=step)  # type: ignore[union-attr]
        except Exception as e:
            logger.warning(f"Failed to log metrics: {e}")

    def log_metric(
        self,
        key: str,
        value: float,
        step: int | None = None,
    ) -> None:
        """Log a single metric.

        Args:
            key: Metric name.
            value: Metric value.
            step: Optional step number.
        """
        self.log_metrics({key: value}, step=step)

    def log_artifact(
        self,
        local_path: str,
        artifact_path: str | None = None,
    ) -> None:
        """Log an artifact file.

        Args:
            local_path: Path to the artifact file.
            artifact_path: Optional subdirectory in artifact store.
        """
        if not self.is_enabled or not self._active_run:
            return

        try:
            mlflow.log_artifact(local_path, artifact_path)  # type: ignore[union-attr]
        except Exception as e:
            logger.warning(f"Failed to log artifact: {e}")

    def log_dict(
        self,
        dictionary: dict[str, Any],
        artifact_file: str,
    ) -> None:
        """Log a dictionary as a JSON artifact.

        Args:
            dictionary: Dictionary to log.
            artifact_file: Name for the artifact file.
        """
        if not self.is_enabled or not self._active_run:
            return

        try:
            mlflow.log_dict(dictionary, artifact_file)  # type: ignore[union-attr]
        except Exception as e:
            logger.warning(f"Failed to log dict: {e}")

    def set_tag(self, key: str, value: str) -> None:
        """Set a tag on the current run.

        Args:
            key: Tag key.
            value: Tag value.
        """
        if not self.is_enabled or not self._active_run:
            return

        try:
            mlflow.set_tag(key, value)  # type: ignore[union-attr]
        except Exception as e:
            logger.warning(f"Failed to set tag: {e}")


@dataclass
class RunContext:
    """Context for an active experiment run."""

    tracker: ExperimentTracker
    run: Any  # mlflow.ActiveRun or None

    @property
    def run_id(self) -> str | None:
        """Get the run ID."""
        if self.run is None:
            return None
        return self.run.info.run_id  # type: ignore[no-any-return]

    def log_params(self, params: dict[str, Any]) -> None:
        """Log parameters."""
        self.tracker.log_params(params)

    def log_metrics(self, metrics: dict[str, float], step: int | None = None) -> None:
        """Log metrics."""
        self.tracker.log_metrics(metrics, step)

    def log_artifact(self, local_path: str, artifact_path: str | None = None) -> None:
        """Log artifact."""
        self.tracker.log_artifact(local_path, artifact_path)


# Global tracker instance
_global_tracker: ExperimentTracker | None = None


def get_tracker(
    experiment_name: str | None = None,
    tracking_uri: str | None = None,
) -> ExperimentTracker:
    """Get or create the global experiment tracker.

    Args:
        experiment_name: Experiment name (used on first call).
        tracking_uri: Tracking URI (used on first call).

    Returns:
        ExperimentTracker: Global tracker instance.
    """
    global _global_tracker
    if _global_tracker is None:
        _global_tracker = ExperimentTracker(
            experiment_name=experiment_name,
            tracking_uri=tracking_uri,
        )
    return _global_tracker


def track_experiment(
    run_name: str | None = None,
    log_params: list[str] | None = None,
    log_result: bool = True,
) -> Callable[[F], F]:
    """Decorator for tracking function execution as an experiment.

    Args:
        run_name: Name for the run (defaults to function name).
        log_params: List of parameter names to log from function arguments.
        log_result: Whether to log the function result.

    Returns:
        Decorated function.

    Usage:
        @track_experiment(run_name="rag-query", log_params=["query", "model"])
        def process_query(query: str, model: str) -> dict:
            ...
    """

    def decorator(func: F) -> F:
        @functools.wraps(func)
        def wrapper(*args: Any, **kwargs: Any) -> Any:
            tracker = get_tracker()
            name = run_name or func.__name__

            with tracker.start_run(run_name=name) as run:
                # Log specified parameters
                if log_params:
                    params_to_log = {}
                    # Get parameter names from function signature
                    import inspect

                    sig = inspect.signature(func)
                    param_names = list(sig.parameters.keys())

                    # Map args to parameter names
                    for i, arg in enumerate(args):
                        if i < len(param_names) and param_names[i] in log_params:
                            params_to_log[param_names[i]] = arg

                    # Add kwargs
                    for key, value in kwargs.items():
                        if key in log_params:
                            params_to_log[key] = value

                    run.log_params(params_to_log)

                # Execute function and track timing
                start_time = time.time()
                try:
                    result = func(*args, **kwargs)
                    tracker.log_metrics({
                        "duration_seconds": time.time() - start_time,
                        "success": 1.0,
                    })

                    if log_result and isinstance(result, dict):
                        # Log numeric values from result
                        metrics = {
                            k: v for k, v in result.items() if isinstance(v, (int, float))
                        }
                        if metrics:
                            tracker.log_metrics(metrics)

                    return result
                except Exception as e:
                    tracker.log_metrics({
                        "duration_seconds": time.time() - start_time,
                        "success": 0.0,
                    })
                    tracker.set_tag("error", str(e))
                    raise

        return wrapper  # type: ignore[return-value]

    return decorator
