"""Calibration artifact persistence.

Handles versioned save/load of fitted calibrator objects alongside
metadata including training provenance, quality metrics, and version.
"""

from __future__ import annotations

import json
import logging
import pickle
from pathlib import Path

from services.intelligence_pipeline_v3.confidence.calibrator import ConfidenceCalibrator
from services.intelligence_pipeline_v3.confidence.models import CalibrationArtifactMetadata

logger = logging.getLogger(__name__)

ARTIFACT_FILE = "calibrator.pkl"
METADATA_FILE = "metadata.json"


def save_artifact(
    calibrator: ConfidenceCalibrator,
    version: str,
    path: str | Path,
) -> Path:
    """Save a fitted calibrator and metadata to a versioned directory.

    Creates the directory structure:
        <path>/<version>/calibrator.pkl
        <path>/<version>/metadata.json

    Parameters
    ----------
    calibrator
        A fitted ConfidenceCalibrator instance.
    version
        Version string for this artifact (e.g., "v1.0.0").
    path
        Base directory for artifact storage.

    Returns
    -------
    Path
        Path to the versioned artifact directory.

    Raises
    ------
    ValueError
        If the calibrator has not been fitted.
    """
    if not calibrator.is_fitted:
        raise ValueError("Cannot save an unfitted calibrator")

    artifact_dir = Path(path) / version
    artifact_dir.mkdir(parents=True, exist_ok=True)

    # Save the calibrator model
    calibrator_path = artifact_dir / ARTIFACT_FILE
    with open(calibrator_path, "wb") as f:
        pickle.dump(calibrator, f, protocol=pickle.HIGHEST_PROTOCOL)

    # Save metadata
    metadata = calibrator.metadata
    if metadata is None:
        metadata = CalibrationArtifactMetadata(
            version=version,
            method=calibrator.method,  # type: ignore[arg-type]
            training_count=0,
            training_range="unknown",
            ece=0.0,
            brier_score=0.0,
        )

    metadata_path = artifact_dir / METADATA_FILE
    with open(metadata_path, "w") as f:
        json.dump(metadata.model_dump(mode="json"), f, indent=2, default=str)

    logger.info(
        "Saved calibration artifact: version=%s, method=%s, path=%s",
        version,
        calibrator.method,
        artifact_dir,
    )
    return artifact_dir


def load_artifact(path: str | Path) -> ConfidenceCalibrator:
    """Load a calibrator from a versioned artifact directory.

    Expects the directory to contain calibrator.pkl and metadata.json.

    Parameters
    ----------
    path
        Path to the versioned artifact directory (e.g., <base>/v1.0.0/).

    Returns
    -------
    ConfidenceCalibrator
        The loaded and ready-to-use calibrator.

    Raises
    ------
    FileNotFoundError
        If the artifact directory or files don't exist.
    ValueError
        If the loaded object is not a ConfidenceCalibrator.
    """
    artifact_dir = Path(path)

    calibrator_path = artifact_dir / ARTIFACT_FILE
    if not calibrator_path.exists():
        raise FileNotFoundError(
            f"Calibrator artifact not found at {calibrator_path}"
        )

    with open(calibrator_path, "rb") as f:
        calibrator = pickle.load(f)  # noqa: S301

    if not isinstance(calibrator, ConfidenceCalibrator):
        raise ValueError(
            f"Loaded object is not a ConfidenceCalibrator: {type(calibrator)}"
        )

    logger.info(
        "Loaded calibration artifact: version=%s, method=%s, path=%s",
        calibrator.version,
        calibrator.method,
        artifact_dir,
    )
    return calibrator


def load_metadata(path: str | Path) -> CalibrationArtifactMetadata:
    """Load only the metadata for an artifact without loading the full model.

    Parameters
    ----------
    path
        Path to the versioned artifact directory.

    Returns
    -------
    CalibrationArtifactMetadata
        The artifact metadata.

    Raises
    ------
    FileNotFoundError
        If the metadata file doesn't exist.
    """
    metadata_path = Path(path) / METADATA_FILE
    if not metadata_path.exists():
        raise FileNotFoundError(f"Metadata not found at {metadata_path}")

    with open(metadata_path) as f:
        data = json.load(f)

    return CalibrationArtifactMetadata(**data)


def list_versions(base_path: str | Path) -> list[str]:
    """List all available artifact versions in a base directory.

    Parameters
    ----------
    base_path
        Base directory containing versioned subdirectories.

    Returns
    -------
    list[str]
        Sorted list of version strings.
    """
    base = Path(base_path)
    if not base.exists():
        return []

    versions = []
    for item in base.iterdir():
        if item.is_dir() and (item / ARTIFACT_FILE).exists():
            versions.append(item.name)

    return sorted(versions)