"""
Portable pickle loader with backwards-compatibility remapping.

Problem
-------
The legacy .pkl was created while the class definitions lived inside
``local-module/train.py``.  Python's pickle protocol stores the fully-
qualified module path (e.g. ``train.AdvancedFeatureEngineer``).  If that
module is not importable, unpickling raises ``ModuleNotFoundError``.

Solution
--------
``_CompatUnpickler`` intercepts every ``find_class`` call and transparently
redirects any legacy path to its canonical location inside this package.

Supported remappings
--------------------
``train``           → ``model_lib._core.architecture``
``simple_predict``  → ``model_lib._core.architecture``
``predict``         → ``model_lib._core.architecture``
``predict_api``     → ``model_lib._core.architecture``
``__main__``        → ``model_lib._core.architecture``  (interactive sessions)
"""
from __future__ import annotations

import io
import os
import pickle
from pathlib import Path
from typing import Any


# ── Module path remapping table ───────────────────────────────────────────────

_LEGACY_MODULE_MAP: dict[str, str] = {
    # bare module names (pkl created while running train.py directly)
    "train": "model_lib._core.architecture",
    "simple_predict": "model_lib._core.architecture",
    "predict": "model_lib._core.architecture",
    "predict_api": "model_lib._core.architecture",
    "__main__": "model_lib._core.architecture",
    # fully-qualified paths when waste_predictor was installed as a package
    "waste_predictor.train": "model_lib._core.architecture",
    "waste_predictor.predict": "model_lib._core.architecture",
    "waste_predictor.predict_api": "model_lib._core.architecture",
    "waste_predictor.simple_predict": "model_lib._core.architecture",
}

# Class-level overrides (module, class) → (new_module, new_class).
# Add entries here if a class was ever renamed between versions.
_LEGACY_CLASS_MAP: dict[tuple[str, str], tuple[str, str]] = {}


class _CompatUnpickler(pickle.Unpickler):
    """pickle.Unpickler that remaps legacy module/class paths on the fly."""

    def find_class(self, module: str, name: str) -> Any:
        # 1. Per-class override
        mapped_cls = _LEGACY_CLASS_MAP.get((module, name))
        if mapped_cls:
            module, name = mapped_cls

        # 2. Module-level remap
        canonical = _LEGACY_MODULE_MAP.get(module, module)

        return super().find_class(canonical, name)


# ── Model discovery ───────────────────────────────────────────────────────────

_DEFAULT_SEARCH_PATHS: list[Path] = [
    # 1. Explicit env-var override
    *(
        [Path(os.environ["WASTE_PREDICTOR_MODEL_PATH"])]
        if "WASTE_PREDICTOR_MODEL_PATH" in os.environ
        else []
    ),
    # 2. Home directory cache
    Path.home() / ".model_lib" / "waste_predictor_v4.pkl",
    # 3. Package-bundled artifacts/
    Path(__file__).parent.parent / "artifacts" / "waste_predictor_v4.pkl",
    # 4. Current working directory (dev convenience)
    Path.cwd() / "waste_predictor_v4.pkl",
]


def _resolve_model_path(explicit: str | None) -> Path:
    """Return a Path to the .pkl file, searching well-known locations."""
    if explicit:
        p = Path(explicit)
        if not p.exists():
            raise FileNotFoundError(f"Explicit model path not found: {p}")
        return p

    for candidate in _DEFAULT_SEARCH_PATHS:
        if candidate.exists():
            return candidate

    searched = "\n  ".join(str(p) for p in _DEFAULT_SEARCH_PATHS)
    raise FileNotFoundError(
        "waste_predictor_v4.pkl not found.\n"
        "Place it in one of:\n"
        f"  {searched}\n"
        "or set the WASTE_PREDICTOR_MODEL_PATH environment variable."
    )


# ── Public loading helpers ────────────────────────────────────────────────────

def load_raw(path: str | Path) -> Any:
    """Deserialise any .pkl produced by local-module, regardless of origin."""
    with open(path, "rb") as f:
        return _CompatUnpickler(f).load()


def load_raw_bytes(data: bytes) -> Any:
    """Deserialise a .pkl from an in-memory bytes object (e.g. from S3 / SQS)."""
    return _CompatUnpickler(io.BytesIO(data)).load()


def download_model_from_s3(
    s3_uri: str,
    destination: str | Path | None = None,
) -> Path:
    """
    Download a model file from S3 and save it locally.

    Parameters
    ----------
    s3_uri:
        S3 URI in the format: s3://bucket-name/path/to/model.pkl
    destination:
        Local path to save the model. If None, saves to the default location:
        ``~/.model_lib/waste_predictor_v4.pkl``

    Returns
    -------
    Path
        The absolute path where the model was saved.

    Raises
    ------
    ImportError:
        If boto3 is not installed (install with: pip install salt_waste_prediction_model[sqs])
    ValueError:
        If the S3 URI format is invalid.

    Examples
    --------
    >>> from model_lib._core.loader import download_model_from_s3
    >>> path = download_model_from_s3("s3://my-bucket/models/waste_predictor_v4.pkl")
    >>> print(f"Model downloaded to: {path}")
    """
    try:
        import boto3
    except ImportError:
        raise ImportError(
            "boto3 is required for S3 downloads. "
            "Install with: pip install salt_waste_prediction_model[sqs]"
        )

    # Parse S3 URI
    if not s3_uri.startswith("s3://"):
        raise ValueError(
            f"Invalid S3 URI: {s3_uri}. Expected format: s3://bucket-name/path/to/file"
        )

    s3_path = s3_uri[5:]  # Remove 's3://'
    parts = s3_path.split("/", 1)
    if len(parts) != 2:
        raise ValueError(
            f"Invalid S3 URI: {s3_uri}. Expected format: s3://bucket-name/path/to/file"
        )

    bucket_name, object_key = parts

    # Determine destination path
    if destination is None:
        destination = Path.home() / ".model_lib" / "waste_predictor_v4.pkl"
    else:
        destination = Path(destination)

    # Create parent directory if it doesn't exist
    destination.parent.mkdir(parents=True, exist_ok=True)

    # Download from S3
    s3_client = boto3.client("s3")
    try:
        s3_client.download_file(bucket_name, object_key, str(destination))
    except Exception as e:
        raise RuntimeError(
            f"Failed to download model from {s3_uri}: {e}"
        ) from e

    return destination


def load_model_file(explicit_path: str | None = None) -> Any:
    """
    Load the ProductionWastePredictor from disk.

    Parameters
    ----------
    explicit_path:
        Absolute path to the .pkl file.  If ``None`` the loader searches the
        default locations defined in ``_DEFAULT_SEARCH_PATHS``.

    Returns
    -------
    ProductionWastePredictor
    """
    path = _resolve_model_path(explicit_path)
    raw = load_raw(path)

    # raw is a dict: {"models": …, "weights": …, "feature_names": …, …}
    # Reconstruct the ProductionWastePredictor wrapper.
    from model_lib._core.architecture import (
        AdvancedFeatureEngineer,
        ProductionWastePredictor,
    )

    predictor = ProductionWastePredictor.__new__(ProductionWastePredictor)
    predictor.models = raw["models"]
    predictor.weights = raw["weights"]
    predictor.output_cols = raw.get("output_cols", AdvancedFeatureEngineer.OUTPUT_COLS)
    predictor.feature_engineer = AdvancedFeatureEngineer()
    predictor.feature_engineer.feature_names = raw.get("feature_names", [])
    predictor.test_metrics = raw.get("metrics", {})
    return predictor
