"""registry.py.

===========
ModelRegistry - helper for managing model lifecycle inside the MLflow
Model Registry (listing, aliasing, loading, cleanup).

Works in concert with :class:`~src.uploader.ModelUploader` and
:class:`~src.naming.ModelNameBuilder`.
"""

from __future__ import annotations

import logging
from typing import List

import mlflow
from mlflow import MlflowClient
from mlflow.entities.model_registry import ModelVersion

from .naming import ModelNameBuilder

logger = logging.getLogger(__name__)

# Alias constants used across the TAT project
ALIAS_CHAMPION = "champion"  # Current production model
ALIAS_CHALLENGER = "challenger"  # A/B test candidate
ALIAS_CANARY = "canary"  # Small-traffic rollout


class ModelRegistry:
    """High-level wrapper around :class:`mlflow.MlflowClient` for the TAT Model Registry.

    Parameters
    ----------
    tracking_uri:
        MLflow tracking server URI. Falls back to the
        ``MLFLOW_TRACKING_URI`` environment variable.

    Examples:
    --------
    >>> from src.registry import ModelRegistry
    >>> reg = ModelRegistry("http://mlflow-server:5000")

    # List all versions of a model
    >>> versions = reg.list_versions("TAT_TwoTower_Baseline")

    # Promote latest version to champion
    >>> reg.promote(model_name="TAT_TwoTower_Baseline", alias=ModelRegistry.ALIAS_CHAMPION)

    # Load champion model for inference
    >>> model = reg.load_champion("TAT_TwoTower_Baseline")
    """

    ALIAS_CHAMPION = ALIAS_CHAMPION
    ALIAS_CHALLENGER = ALIAS_CHALLENGER
    ALIAS_CANARY = ALIAS_CANARY

    def __init__(self, tracking_uri: str | None = None) -> None:
        """Initialize."""
        import os

        uri = tracking_uri or os.environ.get("MLFLOW_TRACKING_URI", "http://127.0.0.1:5001")
        mlflow.set_tracking_uri(uri)
        self._client = MlflowClient()
        self._namer = ModelNameBuilder()
        logger.info("ModelRegistry connected to: %s", uri)

    # ------------------------------------------------------------------
    # Read operations
    # ------------------------------------------------------------------

    def list_versions(self, model_name: str) -> List[ModelVersion]:
        """Return all registered versions of *model_name*, newest first.

        Parameters
        ----------
        model_name:
            Registered model name in MLflow (e.g. ``"TAT_TwoTower_Baseline"``).
            Use :meth:`~naming.ModelNameBuilder.registered_model_name` to
            generate the correct name.

        Returns:
        -------
        list[ModelVersion]
        """
        versions: List[ModelVersion] = self._client.search_model_versions(f"name='{model_name}'")
        versions_sorted = sorted(versions, key=lambda v: int(v.version), reverse=True)
        logger.info("Found %d version(s) for model '%s'.", len(versions_sorted), model_name)
        return versions_sorted

    def get_latest_version(self, model_name: str) -> ModelVersion | None:
        """Return the highest-numbered version of *model_name*, or ``None``."""
        versions = self.list_versions(model_name)
        return versions[0] if versions else None

    def get_version_by_alias(self, model_name: str, alias: str) -> ModelVersion | None:
        """Return the :class:`ModelVersion` currently assigned to *alias*.

        Parameters
        ----------
        model_name:
            Registered model name.
        alias:
            Alias string (e.g. ``"champion"``).

        Returns:
        -------
        ModelVersion or None
        """
        try:
            mv = self._client.get_model_version_by_alias(model_name, alias)
            return mv
        except mlflow.exceptions.MlflowException:
            logger.warning("Alias '%s' not found for model '%s'.", alias, model_name)
            return None

    # ------------------------------------------------------------------
    # Write / lifecycle operations
    # ------------------------------------------------------------------

    def promote(
        self,
        model_name: str,
        alias: str = ALIAS_CHAMPION,
        version: str | None = None,
    ) -> ModelVersion:
        """Assign *alias* to a model version.

        Parameters
        ----------
        model_name:
            Registered model name.
        alias:
            Alias to set (default: ``"champion"``).
        version:
            Version number string (e.g. ``"3"``). If omitted, the latest
            version is used.

        Returns:
        -------
        ModelVersion
            The version that was aliased.

        Raises:
        ------
        RuntimeError
            If no versions exist for the model.
        """
        if version is None:
            latest = self.get_latest_version(model_name)
            if latest is None:
                raise RuntimeError(f"No versions found for '{model_name}'. Upload a model first.")
            version = latest.version

        self._client.set_registered_model_alias(model_name, alias, version)
        mv = self._client.get_model_version(model_name, version)

        logger.info("Set alias '%s' → '%s' version %s.", alias, model_name, version)
        return mv

    def set_tag(self, model_name: str, version: str, key: str, value: str) -> None:
        """Add or update a tag on a specific model version."""
        self._client.set_model_version_tag(model_name, version, key, value)
        logger.info("Tagged '%s' v%s : %s=%s", model_name, version, key, value)

    def delete_version(self, model_name: str, version: str) -> None:
        """Permanently delete a model version from the registry.

        .. warning::
            This action is irreversible.
        """
        self._client.delete_model_version(model_name, version)
        logger.warning("DELETED '%s' version %s from registry.", model_name, version)

    def cleanup_old_versions(
        self,
        model_name: str,
        keep: int = 5,
        dry_run: bool = True,
    ) -> List[str]:
        """Remove all but the *keep* newest versions of *model_name*.

        Parameters
        ----------
        model_name:
            Registered model name.
        keep:
            Number of most recent versions to retain.
        dry_run:
            If ``True`` (default), log which versions *would* be deleted
            but do not actually delete them.

        Returns:
        -------
        list[str]
            Version numbers that were (or would be) deleted.
        """
        versions = self.list_versions(model_name)
        to_delete = versions[keep:]

        deleted: List[str] = []
        for mv in to_delete:
            logger.info(
                "[%s] Would delete '%s' version %s (run_id=%s).",
                "DRY-RUN" if dry_run else "DELETING",
                model_name,
                mv.version,
                mv.run_id,
            )
            if not dry_run:
                self.delete_version(model_name, mv.version)
            deleted.append(mv.version)

        return deleted

    # ------------------------------------------------------------------
    # Load operations
    # ------------------------------------------------------------------

    def load_champion(self, model_name: str) -> mlflow.pyfunc.PyFuncModel:
        """Load the model currently aliased as ``"champion"``.

        Parameters
        ----------
        model_name:
            Registered model name.

        Returns:
        -------
        PyFuncModel
        """
        uri = f"models:/{model_name}@{ALIAS_CHAMPION}"
        logger.info("Loading champion model from: %s", uri)
        return mlflow.pyfunc.load_model(uri)

    def load_version(self, model_name: str, version: str) -> mlflow.pyfunc.PyFuncModel:
        """Load a specific model version by number.

        Parameters
        ----------
        model_name:
            Registered model name.
        version:
            Version number string (e.g. ``"2"``).

        Returns:
        -------
        PyFuncModel
        """
        uri = f"models:/{model_name}/{version}"
        logger.info("Loading model from: %s", uri)
        return mlflow.pyfunc.load_model(uri)

    # ------------------------------------------------------------------
    # Reporting
    # ------------------------------------------------------------------

    def print_summary(self, model_name: str) -> None:
        """Print a formatted summary of all versions for *model_name*."""
        versions = self.list_versions(model_name)

        print(f"\n{'=' * 64}")
        print(f"  Model Registry Summary: {model_name}")
        print(f"{'=' * 64}")
        print(f"  {'Version':<10} {'Run ID':<36} {'Last Updated'}")
        print(f"  {'-' * 60}")
        for mv in versions:
            ts = mv.last_updated_timestamp
            import datetime

            dt_str = datetime.datetime.fromtimestamp(ts / 1000).strftime("%Y-%m-%d %H:%M") if ts else "N/A"
            print(f"  v{mv.version:<9} {mv.run_id:<36} {dt_str}")

        print(f"{'=' * 64}\n")
