from typing import Optional, Dict, Any
import mlflow
from mlflow.tracking import MlflowClient
from mloptiflow.utils.exceptions import PredictionError
import os
from pathlib import Path
import logging


class ModelRegistry:
    def __init__(self, tracking_uri: Optional[str] = None):
        if tracking_uri:
            mlflow.set_tracking_uri(tracking_uri)
            logging.debug(f"Using explicit tracking URI: {tracking_uri}")
        else:
            env_uri = os.getenv("MLFLOW_TRACKING_URI", "file:./mlruns")
            mlflow.set_tracking_uri(env_uri)
            logging.debug(f"Using environment tracking URI: {env_uri}")

        self.client = MlflowClient()
        logging.debug(f"Final tracking URI: {mlflow.get_tracking_uri()}")
        self._validate_registry_path()

    def _validate_registry_path(self):
        tracking_uri = mlflow.get_tracking_uri()
        logging.debug(f"Validating registry path: {tracking_uri}")
        if tracking_uri.startswith("file://"):
            path = Path(tracking_uri[7:])
            logging.debug(f"Resolved local path: {path.absolute()}")
            if not path.exists():
                raise ValueError(f"MLflow registry path {path} does not exist")

    def get_latest_model(
        self, experiment_name: str, model_name: str = "XGBoost"
    ) -> Dict[str, Any]:
        experiment = self.client.get_experiment_by_name(experiment_name)
        if experiment is None:
            raise PredictionError(f"No experiment found with name: {experiment_name}")

        runs = self.client.search_runs(
            experiment_ids=[experiment.experiment_id],
            filter_string="status = 'FINISHED'",
            order_by=["start_time DESC"],
            max_results=1,
        )

        if not runs:
            raise PredictionError("No successful runs found")

        run = runs[0]

        return {
            "run_id": run.info.run_id,
            "model_uri": f"runs:/{run.info.run_id}/{model_name}_model",
            "scaler_uri": f"runs:/{run.info.run_id}/scaler",
        }

    def load_model(
        self, experiment_name: str, model_name: str = "XGBoost"
    ) -> Dict[str, Any]:
        model_info = self.get_latest_model(experiment_name, model_name)
        return {
            "model": mlflow.sklearn.load_model(model_info["model_uri"]),
            "scaler": mlflow.sklearn.load_model(model_info["scaler_uri"]),
        }
