"""
Experiment tracking and MLOps integration.

Supports:
- MLflow for experiment tracking
- Weights & Biases (W&B) integration
- Local file-based tracking as fallback
"""

import json
import os
import time
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional, Union

from ..infra.logging import get_logger

logger = get_logger(__name__)


@dataclass
class ExperimentConfig:
    """Configuration for an experiment run."""
    experiment_name: str
    run_name: Optional[str] = None
    tags: Dict[str, str] = field(default_factory=dict)
    
    # Model configuration
    model_name: str = ""
    model_version: str = ""
    
    # RAG configuration
    embedding_model: str = ""
    chunk_size: int = 0
    chunk_overlap: int = 0
    retrieval_top_k: int = 0
    
    # Prompt configuration
    prompt_version: str = ""
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            "experiment_name": self.experiment_name,
            "run_name": self.run_name,
            "tags": self.tags,
            "model_name": self.model_name,
            "model_version": self.model_version,
            "embedding_model": self.embedding_model,
            "chunk_size": self.chunk_size,
            "chunk_overlap": self.chunk_overlap,
            "retrieval_top_k": self.retrieval_top_k,
            "prompt_version": self.prompt_version,
        }


@dataclass
class ExperimentMetrics:
    """Metrics from an experiment run."""
    # Quality metrics
    concept_count: int = 0
    formula_count: int = 0
    algorithm_count: int = 0
    
    # Coverage metrics
    topic_coverage: float = 0.0
    
    # Performance metrics
    total_duration_ms: float = 0.0
    total_tokens: int = 0
    
    # Pass-specific metrics
    pass_metrics: Dict[str, Dict[str, Any]] = field(default_factory=dict)
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            "concept_count": self.concept_count,
            "formula_count": self.formula_count,
            "algorithm_count": self.algorithm_count,
            "topic_coverage": self.topic_coverage,
            "total_duration_ms": self.total_duration_ms,
            "total_tokens": self.total_tokens,
            "pass_metrics": self.pass_metrics,
        }


class ExperimentTracker(ABC):
    """Abstract base class for experiment tracking."""
    
    @abstractmethod
    def start_run(self, config: ExperimentConfig) -> str:
        """Start a new experiment run. Returns run ID."""
        pass
    
    @abstractmethod
    def log_params(self, params: Dict[str, Any]) -> None:
        """Log parameters for the current run."""
        pass
    
    @abstractmethod
    def log_metrics(self, metrics: Dict[str, float], step: Optional[int] = None) -> None:
        """Log metrics for the current run."""
        pass
    
    @abstractmethod
    def log_artifact(self, local_path: str, artifact_name: Optional[str] = None) -> None:
        """Log an artifact (file) for the current run."""
        pass
    
    @abstractmethod
    def end_run(self, status: str = "completed") -> None:
        """End the current run."""
        pass
    
    @abstractmethod
    def get_run_id(self) -> Optional[str]:
        """Get current run ID."""
        pass


class MLflowTracker(ExperimentTracker):
    """MLflow-based experiment tracker."""
    
    def __init__(
        self,
        tracking_uri: Optional[str] = None,
        artifact_location: Optional[str] = None
    ):
        """
        Initialize MLflow tracker.
        
        Args:
            tracking_uri: MLflow tracking server URI
            artifact_location: Location for storing artifacts
        """
        self._run_id: Optional[str] = None
        self._mlflow = None
        
        try:
            import mlflow
            self._mlflow = mlflow
            
            if tracking_uri:
                mlflow.set_tracking_uri(tracking_uri)
            
            logger.info(f"MLflow tracker initialized, tracking URI: {mlflow.get_tracking_uri()}")
            
        except ImportError:
            logger.warning("MLflow not installed. Install with: pip install mlflow")
            raise ImportError("mlflow is required for MLflowTracker")
    
    def start_run(self, config: ExperimentConfig) -> str:
        """Start a new MLflow run."""
        self._mlflow.set_experiment(config.experiment_name)
        
        run = self._mlflow.start_run(run_name=config.run_name, tags=config.tags)
        self._run_id = run.info.run_id
        
        # Log config as params
        self.log_params(config.to_dict())
        
        logger.info(f"Started MLflow run: {self._run_id}")
        return self._run_id
    
    def log_params(self, params: Dict[str, Any]) -> None:
        """Log parameters to MLflow."""
        for key, value in params.items():
            if isinstance(value, dict):
                for k, v in value.items():
                    self._mlflow.log_param(f"{key}.{k}", v)
            else:
                self._mlflow.log_param(key, value)
    
    def log_metrics(self, metrics: Dict[str, float], step: Optional[int] = None) -> None:
        """Log metrics to MLflow."""
        self._mlflow.log_metrics(metrics, step=step)
    
    def log_artifact(self, local_path: str, artifact_name: Optional[str] = None) -> None:
        """Log artifact to MLflow."""
        self._mlflow.log_artifact(local_path, artifact_name)
    
    def end_run(self, status: str = "completed") -> None:
        """End MLflow run."""
        mlflow_status = "FINISHED" if status == "completed" else "FAILED"
        self._mlflow.end_run(status=mlflow_status)
        logger.info(f"Ended MLflow run: {self._run_id}")
        self._run_id = None
    
    def get_run_id(self) -> Optional[str]:
        """Get current run ID."""
        return self._run_id


class WandbTracker(ExperimentTracker):
    """Weights & Biases experiment tracker."""
    
    def __init__(
        self,
        project: str,
        entity: Optional[str] = None
    ):
        """
        Initialize W&B tracker.
        
        Args:
            project: W&B project name
            entity: W&B entity (team/user)
        """
        self._run = None
        self._wandb = None
        self._project = project
        self._entity = entity
        
        try:
            import wandb
            self._wandb = wandb
            logger.info(f"W&B tracker initialized for project: {project}")
            
        except ImportError:
            logger.warning("wandb not installed. Install with: pip install wandb")
            raise ImportError("wandb is required for WandbTracker")
    
    def start_run(self, config: ExperimentConfig) -> str:
        """Start a new W&B run."""
        self._run = self._wandb.init(
            project=self._project,
            entity=self._entity,
            name=config.run_name,
            tags=list(config.tags.values()),
            config=config.to_dict()
        )
        
        logger.info(f"Started W&B run: {self._run.id}")
        return self._run.id
    
    def log_params(self, params: Dict[str, Any]) -> None:
        """Log parameters to W&B."""
        if self._run:
            self._run.config.update(params)
    
    def log_metrics(self, metrics: Dict[str, float], step: Optional[int] = None) -> None:
        """Log metrics to W&B."""
        if self._run:
            self._wandb.log(metrics, step=step)
    
    def log_artifact(self, local_path: str, artifact_name: Optional[str] = None) -> None:
        """Log artifact to W&B."""
        if self._run:
            artifact = self._wandb.Artifact(
                artifact_name or Path(local_path).stem,
                type="output"
            )
            artifact.add_file(local_path)
            self._run.log_artifact(artifact)
    
    def end_run(self, status: str = "completed") -> None:
        """End W&B run."""
        if self._run:
            self._run.finish()
            logger.info(f"Ended W&B run: {self._run.id}")
            self._run = None
    
    def get_run_id(self) -> Optional[str]:
        """Get current run ID."""
        return self._run.id if self._run else None


class LocalFileTracker(ExperimentTracker):
    """
    Local file-based experiment tracker.
    
    Fallback when MLflow/W&B are not available.
    """
    
    def __init__(self, output_dir: Union[str, Path] = "./experiments"):
        """
        Initialize local tracker.
        
        Args:
            output_dir: Directory for storing experiment data
        """
        self._output_dir = Path(output_dir)
        self._output_dir.mkdir(parents=True, exist_ok=True)
        self._run_id: Optional[str] = None
        self._run_dir: Optional[Path] = None
        self._data: Dict[str, Any] = {}
        
        logger.info(f"Local tracker initialized, output dir: {self._output_dir}")
    
    def start_run(self, config: ExperimentConfig) -> str:
        """Start a new local run."""
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        self._run_id = f"{config.experiment_name}_{timestamp}"
        self._run_dir = self._output_dir / self._run_id
        self._run_dir.mkdir(parents=True, exist_ok=True)
        
        self._data = {
            "run_id": self._run_id,
            "start_time": datetime.now().isoformat(),
            "config": config.to_dict(),
            "params": {},
            "metrics": [],
            "artifacts": [],
            "status": "running"
        }
        
        self._save_data()
        logger.info(f"Started local run: {self._run_id}")
        return self._run_id
    
    def log_params(self, params: Dict[str, Any]) -> None:
        """Log parameters locally."""
        self._data["params"].update(params)
        self._save_data()
    
    def log_metrics(self, metrics: Dict[str, float], step: Optional[int] = None) -> None:
        """Log metrics locally."""
        entry = {
            "timestamp": datetime.now().isoformat(),
            "step": step,
            "metrics": metrics
        }
        self._data["metrics"].append(entry)
        self._save_data()
    
    def log_artifact(self, local_path: str, artifact_name: Optional[str] = None) -> None:
        """Copy artifact to run directory."""
        import shutil
        
        src = Path(local_path)
        if src.exists() and self._run_dir:
            dst = self._run_dir / (artifact_name or src.name)
            shutil.copy2(src, dst)
            self._data["artifacts"].append(str(dst))
            self._save_data()
    
    def end_run(self, status: str = "completed") -> None:
        """End local run."""
        self._data["end_time"] = datetime.now().isoformat()
        self._data["status"] = status
        self._save_data()
        logger.info(f"Ended local run: {self._run_id}")
        self._run_id = None
        self._run_dir = None
    
    def get_run_id(self) -> Optional[str]:
        """Get current run ID."""
        return self._run_id
    
    def _save_data(self) -> None:
        """Save run data to JSON file."""
        if self._run_dir:
            with open(self._run_dir / "run_data.json", "w") as f:
                json.dump(self._data, f, indent=2)


def create_tracker(
    backend: str = "local",
    **kwargs
) -> ExperimentTracker:
    """
    Factory function to create experiment tracker.
    
    Args:
        backend: Tracker backend ("mlflow", "wandb", "local")
        **kwargs: Backend-specific arguments
        
    Returns:
        ExperimentTracker instance
    """
    if backend == "mlflow":
        try:
            return MLflowTracker(
                tracking_uri=kwargs.get("tracking_uri"),
                artifact_location=kwargs.get("artifact_location")
            )
        except ImportError:
            logger.warning("MLflow not available, falling back to local tracker")
            return LocalFileTracker(kwargs.get("output_dir", "./experiments"))
    
    elif backend == "wandb":
        try:
            return WandbTracker(
                project=kwargs.get("project", "readanybook"),
                entity=kwargs.get("entity")
            )
        except ImportError:
            logger.warning("W&B not available, falling back to local tracker")
            return LocalFileTracker(kwargs.get("output_dir", "./experiments"))
    
    else:
        return LocalFileTracker(kwargs.get("output_dir", "./experiments"))
