"""uploader.py.

===========
ModelUploader - generic, framework-agnostic model uploader for MLflow.

Supported upload strategies
----------------------------
* **pyfunc**      - custom :class:`mlflow.pyfunc.PythonModel` wrapper (default)
* **sklearn**     - scikit-learn pipelines / estimators
* **pytorch**     - raw PyTorch ``nn.Module`` state dicts
* **tensorflow**  - TensorFlow / Keras models
* **artifact**    - bag-of-files upload (no registered flavour needed)

All strategies share the same naming convention through
:class:`~src.naming.ModelNameBuilder`.
"""

from __future__ import annotations

import logging
import os
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Dict, List

import mlflow
import mlflow.pyfunc
from mlflow.models import infer_signature

from .naming import (
    Architecture,
    ModelNameBuilder,
    ModelVariant,
    RunAction,
    Stage,
)

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Data-class: upload configuration
# ---------------------------------------------------------------------------


@dataclass
class UploadConfig:
    """All parameters needed to perform a single model upload to MLflow.

    Parameters
    ----------
    architecture:
        Model architecture used for the registered model name.
    variant:
        Model variant (Baseline / Champion / …).
    component:
        Experiment component label (e.g. ``"Recommendation"``).
    stage:
        Experiment lifecycle stage (e.g. ``Stage.RESEARCH``).
    use_case:
        Short label for the run name (e.g. ``"MiniLM"``).
    action:
        Run-name action prefix (e.g. ``RunAction.UPLOAD``).
    tracking_uri:
        MLflow tracking server URI. Falls back to the
        ``MLFLOW_TRACKING_URI`` environment variable, or ``http://127.0.0.1:5001``.
    tags:
        Extra key-value tags logged to the MLflow run.
    params:
        Extra key-value parameters logged to the MLflow run.
    register_model:
        Whether to register the model in the MLflow Model Registry.
    artifact_path:
        Path inside the run where the model artefact is stored (default ``"model"``).
    run_description:
        Human-readable description of the MLflow run.
    """

    architecture: Architecture | str = Architecture.TWO_TOWER
    variant: ModelVariant | str = ModelVariant.BASELINE
    component: str = "Recommendation"
    stage: Stage | str = Stage.RESEARCH
    use_case: str = "Model"
    action: RunAction | str = RunAction.UPLOAD
    tracking_uri: str | None = None
    tags: Dict[str, str] = field(default_factory=dict)
    params: Dict[str, Any] = field(default_factory=dict)
    register_model: bool = True
    artifact_path: str = "model"
    run_description: str = ""

    # Resolved names (populated lazily via ModelNameBuilder)
    _experiment_name: str = field(init=False, default="")
    _registered_model_name: str = field(init=False, default="")
    _run_name: str = field(init=False, default="")

    def __post_init__(self) -> None:
        """Post-init."""
        namer = ModelNameBuilder()
        self._experiment_name = namer.experiment_name(component=self.component, stage=self.stage)
        self._registered_model_name = namer.registered_model_name(architecture=self.architecture, variant=self.variant)
        self._run_name = namer.run_name(action=self.action, use_case=self.use_case)

    # ------------------------------------------------------------------
    # Properties
    # ------------------------------------------------------------------

    @property
    def experiment_name(self) -> str:
        return self._experiment_name

    @property
    def registered_model_name(self) -> str:
        return self._registered_model_name

    @property
    def run_name(self) -> str:
        return self._run_name

    @property
    def effective_tracking_uri(self) -> str:
        return self.tracking_uri or os.environ.get("MLFLOW_TRACKING_URI", "http://127.0.0.1:5001")


# ---------------------------------------------------------------------------
# ModelUploader
# ---------------------------------------------------------------------------


class ModelUploader:
    """Uploads a trained model to MLflow with standard TAT naming conventions.

    Examples:
    --------
    **Upload a custom PythonModel wrapper:**

    >>> from src.uploader import ModelUploader, UploadConfig
    >>> from src.naming import Architecture, ModelVariant, Stage, RunAction

    >>> config = UploadConfig(
    ...     architecture=Architecture.TWO_TOWER,
    ...     variant=ModelVariant.BASELINE,
    ...     component="Recommendation",
    ...     stage=Stage.RESEARCH,
    ...     use_case="MiniLM",
    ...     action=RunAction.UPLOAD,
    ...     params={"embedding_dim": 256},
    ...     tags={"developer": "ai-team"},
    ... )

    >>> uploader = ModelUploader(config)
    >>> result = uploader.upload_pyfunc(
    ...     python_model=my_wrapper,
    ...     artifacts={
    ...         "model_weights": "/path/to/model.pt",
    ...         "item_vectors": "/path/to/vectors.npy",
    ...     },
    ...     signature=my_signature,
    ...     code_paths=["src/model_definitions.py"],
    ... )
    >>> print(result["model_uri"])

    **Upload a scikit-learn model:**

    >>> result = uploader.upload_sklearn(sk_model=pipeline, input_example=X_test[:5])

    **Upload raw artifact files only:**

    >>> result = uploader.upload_artifacts(files=["/path/to/model.pkl"])
    """

    def __init__(self, config: UploadConfig) -> None:
        """Initialize."""
        self.config = config
        self._setup_mlflow()

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _setup_mlflow(self) -> None:
        uri = self.config.effective_tracking_uri
        mlflow.set_tracking_uri(uri)
        mlflow.set_experiment(self.config.experiment_name)
        logger.info("MLflow tracking URI : %s", uri)
        logger.info("MLflow experiment   : %s", self.config.experiment_name)

    def _start_run(self) -> mlflow.ActiveRun:
        return mlflow.start_run(
            run_name=self.config.run_name,
            description=self.config.run_description
            or (f"Upload of {self.config.registered_model_name} ({self.config.variant}) to MLflow registry."),
        )

    def _log_common(self) -> None:
        """Log tags and params that are common to every upload strategy."""
        default_tags = {
            "project": "TAT-Personalization",
            "architecture": str(self.config.architecture),
            "variant": str(self.config.variant),
            "registry_module_version": "0.1.0",
        }
        default_tags.update(self.config.tags)
        mlflow.set_tags(default_tags)

        if self.config.params:
            mlflow.log_params(self.config.params)

    # ------------------------------------------------------------------
    # Public upload methods
    # ------------------------------------------------------------------

    def upload_pyfunc(
        self,
        python_model: mlflow.pyfunc.PythonModel,
        artifacts: Dict[str, str] | None = None,
        signature: Any | None = None,
        input_example: Any | None = None,
        code_paths: List[str] | None = None,
        conda_env: Dict | None = None,
        pip_requirements: List[str] | None = None,
    ) -> Dict[str, Any]:
        """Upload a :class:`mlflow.pyfunc.PythonModel` to MLflow.

        Parameters
        ----------
        python_model:
            Instance of a class that inherits from
            :class:`mlflow.pyfunc.PythonModel`.
        artifacts:
            Mapping of ``{key: local_path}`` that the wrapper's
            ``load_context`` will receive.
        signature:
            :class:`mlflow.models.ModelSignature` describing input/output
            schema. Use :func:`mlflow.models.infer_signature` to create one.
        input_example:
            Representative input sample for documentation.
        code_paths:
            List of Python file/directory paths to bundle with the model
            (e.g. ``["src/model_definitions.py"]``).
        conda_env / pip_requirements:
            Environment specification.

        Returns:
        -------
        dict
            ``{"run_id", "run_name", "model_uri", "registered_model_name"}``.
        """
        registered_name = self.config.registered_model_name if self.config.register_model else None

        logger.info("Starting pyfunc upload → registered name: %s", registered_name)

        if signature is None and input_example is None:
            logger.warning(
                "⚠️ BEST PRACTICE: 'signature' and 'input_example' are missing. "
                "Always log a model signature (schema) to enforce the data contract!"
            )

        with self._start_run() as run:
            self._log_common()

            log_kwargs: Dict[str, Any] = dict(
                name=self.config.artifact_path,
                python_model=python_model,
            )
            if artifacts:
                log_kwargs["artifacts"] = artifacts
            if signature is not None:
                log_kwargs["signature"] = signature
            if input_example is not None:
                log_kwargs["input_example"] = input_example
            if code_paths:
                log_kwargs["code_paths"] = code_paths
            if conda_env:
                log_kwargs["conda_env"] = conda_env
            if pip_requirements:
                log_kwargs["pip_requirements"] = pip_requirements
            if registered_name:
                log_kwargs["registered_model_name"] = registered_name

            model_info = mlflow.pyfunc.log_model(**log_kwargs)

            result = {
                "run_id": run.info.run_id,
                "run_name": self.config.run_name,
                "model_uri": model_info.model_uri,
                "registered_model_name": registered_name,
            }

        logger.info("Upload complete. model_uri=%s", result["model_uri"])
        return result

    def upload_sklearn(
        self,
        sk_model: Any,
        signature: Any | None = None,
        input_example: Any | None = None,
        serialization_format: str = "cloudpickle",
        pip_requirements: List[str] | None = None,
    ) -> Dict[str, Any]:
        """Upload a scikit-learn model / pipeline to MLflow.

        Parameters
        ----------
        sk_model:
            Fitted scikit-learn estimator or pipeline.
        signature:
            Model signature. Auto-inferred from ``input_example`` if omitted.
        input_example:
            Representative input sample  - also used to infer the signature.
        serialization_format:
            ``"cloudpickle"`` (default) or ``"pickle"``.
        pip_requirements:
            Extra Python package requirements.

        Returns:
        -------
        dict
        """
        import mlflow.sklearn

        registered_name = self.config.registered_model_name if self.config.register_model else None

        if signature is None and input_example is not None:
            import pandas as pd

            example_df = input_example if isinstance(input_example, pd.DataFrame) else pd.DataFrame(input_example)
            predictions = sk_model.predict(example_df)
            signature = infer_signature(example_df, predictions)

        if signature is None and input_example is None:
            logger.warning(
                "⚠️ BEST PRACTICE: 'signature' and 'input_example' are missing. "
                "Always log a model signature (schema) to enforce the data contract!"
            )

        logger.info("Starting sklearn upload → registered name: %s", registered_name)

        with self._start_run() as run:
            self._log_common()

            log_kwargs: Dict[str, Any] = dict(
                sk_model=sk_model,
                name=self.config.artifact_path,
                serialization_format=serialization_format,
            )
            if signature is not None:
                log_kwargs["signature"] = signature
            if input_example is not None:
                log_kwargs["input_example"] = input_example
            if pip_requirements:
                log_kwargs["pip_requirements"] = pip_requirements
            if registered_name:
                log_kwargs["registered_model_name"] = registered_name

            model_info = mlflow.sklearn.log_model(**log_kwargs)

            result = {
                "run_id": run.info.run_id,
                "run_name": self.config.run_name,
                "model_uri": model_info.model_uri,
                "registered_model_name": registered_name,
            }

        logger.info("Upload complete. model_uri=%s", result["model_uri"])
        return result

    def upload_pytorch(
        self,
        pytorch_model: Any,
        signature: Any | None = None,
        input_example: Any | None = None,
        extra_files: List[str] | None = None,
        pip_requirements: List[str] | None = None,
    ) -> Dict[str, Any]:
        """Upload a raw PyTorch ``nn.Module`` to MLflow.

        Parameters
        ----------
        pytorch_model:
            Trained :class:`torch.nn.Module` instance.
        signature:
            Model signature.
        input_example:
            Representative input sample.
        extra_files:
            Additional files to include in the MLflow model artefact.
        pip_requirements:
            Extra Python package requirements.

        Returns:
        -------
        dict
        """
        import mlflow.pytorch

        registered_name = self.config.registered_model_name if self.config.register_model else None

        logger.info("Starting pytorch upload → registered name: %s", registered_name)

        if signature is None and input_example is None:
            logger.warning(
                "⚠️ BEST PRACTICE: 'signature' and 'input_example' are missing. "
                "Always log a model signature (schema) to enforce the data contract!"
            )

        with self._start_run() as run:
            self._log_common()

            log_kwargs: Dict[str, Any] = dict(
                pytorch_model=pytorch_model,
                name=self.config.artifact_path,
            )
            if signature is not None:
                log_kwargs["signature"] = signature
            if input_example is not None:
                log_kwargs["input_example"] = input_example
            if extra_files:
                log_kwargs["extra_files"] = extra_files
            if pip_requirements:
                log_kwargs["pip_requirements"] = pip_requirements
            if registered_name:
                log_kwargs["registered_model_name"] = registered_name

            model_info = mlflow.pytorch.log_model(**log_kwargs)

            result = {
                "run_id": run.info.run_id,
                "run_name": self.config.run_name,
                "model_uri": model_info.model_uri,
                "registered_model_name": registered_name,
            }

        logger.info("Upload complete. model_uri=%s", result["model_uri"])
        return result

    def upload_tensorflow(
        self,
        tf_model: Any,
        signature: Any | None = None,
        input_example: Any | None = None,
        pip_requirements: List[str] | None = None,
    ) -> Dict[str, Any]:
        """Upload a raw TensorFlow/Keras model to MLflow.

        Parameters
        ----------
        tf_model:
            Trained TensorFlow or Keras model instance.
        signature:
            Model signature.
        input_example:
            Representative input sample.
        pip_requirements:
            Extra Python package requirements.

        Returns:
        -------
        dict
        """
        import mlflow.tensorflow

        registered_name = self.config.registered_model_name if self.config.register_model else None

        logger.info("Starting tensorflow upload → registered name: %s", registered_name)

        if signature is None and input_example is None:
            logger.warning(
                "⚠️ BEST PRACTICE: 'signature' and 'input_example' are missing. "
                "Always log a model signature (schema) to enforce the data contract!"
            )

        with self._start_run() as run:
            self._log_common()

            log_kwargs: Dict[str, Any] = dict(
                model=tf_model,
                name=self.config.artifact_path,
            )
            if signature is not None:
                log_kwargs["signature"] = signature
            if input_example is not None:
                log_kwargs["input_example"] = input_example
            if pip_requirements:
                log_kwargs["pip_requirements"] = pip_requirements
            if registered_name:
                log_kwargs["registered_model_name"] = registered_name

            model_info = mlflow.tensorflow.log_model(**log_kwargs)

            result = {
                "run_id": run.info.run_id,
                "run_name": self.config.run_name,
                "model_uri": model_info.model_uri,
                "registered_model_name": registered_name,
            }

        logger.info("Upload complete. model_uri=%s", result["model_uri"])
        return result

    def upload_artifacts(
        self,
        files: List[str | Path],
        artifact_subdir: str = "",
        extra_metrics: Dict[str, float] | None = None,
    ) -> Dict[str, Any]:
        """Upload arbitrary files as MLflow artifacts (no model flavour).

        Useful for uploading pre-built artefacts (vectors, configs, CSVs)
        that don't need a registered model flavour.

        Parameters
        ----------
        files:
            List of local file paths to upload.
        artifact_subdir:
            Optional subdirectory inside the run's artifact store.
        extra_metrics:
            Key-value metrics to log alongside the artefacts.

        Returns:
        -------
        dict
            ``{"run_id", "run_name", "artifact_uris"}``.
        """
        logger.info("Uploading %d artifact file(s).", len(files))

        with self._start_run() as run:
            self._log_common()

            if extra_metrics:
                mlflow.log_metrics(extra_metrics)

            artifact_uris: List[str] = []
            for f in files:
                path = Path(f)
                if not path.exists():
                    raise FileNotFoundError(f"Artifact not found: {path}. Ensure the file exists before uploading.")
                mlflow.log_artifact(str(path), artifact_path=artifact_subdir or None)
                artifact_uris.append(str(path.name))
                logger.info("  ✓ Uploaded: %s", path.name)

            result = {
                "run_id": run.info.run_id,
                "run_name": self.config.run_name,
                "artifact_uris": artifact_uris,
            }

        logger.info("Artifact upload complete. run_id=%s", result["run_id"])
        return result

    # ------------------------------------------------------------------
    # Pretty summary
    # ------------------------------------------------------------------

    def summary(self) -> str:
        """Return a human-readable summary of the resolved naming."""
        lines = [
            "=" * 60,
            "  MLflow Upload Configuration",
            "=" * 60,
            f"  Tracking URI      : {self.config.effective_tracking_uri}",
            f"  Experiment        : {self.config.experiment_name}",
            f"  Registered Model  : {self.config.registered_model_name}",
            f"  Run Name          : {self.config.run_name}",
            f"  Artifact Path     : {self.config.artifact_path}",
            f"  Register Model    : {self.config.register_model}",
            "=" * 60,
        ]
        return "\n".join(lines)
