"""Weight registry for fine-tuned ModernBERT LoRA adapters.

Maintains a JSON manifest at ``~/.llmhosts/weights/modernbert/manifest.json``
tracking all registered adapter versions, their metrics, and which version
is currently active.

Version format: ``v{YYYYMMDD}-{n_samples}k``
Example:       ``v20260224-10k``
"""

from __future__ import annotations

import dataclasses
import datetime
import json
import logging
import re
from pathlib import Path
from typing import TYPE_CHECKING, Any

logger = logging.getLogger(__name__)

if TYPE_CHECKING:
    from llmhosts.training.trainer import TrainingResult

_DEFAULT_REGISTRY_DIR: Path = Path.home() / ".llmhosts" / "weights" / "modernbert"
_MANIFEST_FILENAME: str = "manifest.json"


@dataclasses.dataclass
class WeightManifest:
    """Metadata record for a single registered adapter version."""

    version: str
    """Version string, e.g. ``v20260224-10k``."""

    path: Path
    """Absolute path to the LoRA adapter directory."""

    accuracy: float
    """Evaluation accuracy (0.0-1.0) of this adapter."""

    improvement_pp: float
    """Accuracy improvement in percentage points vs baseline."""

    n_training_samples: int
    """Number of training examples used to produce this adapter."""

    created_at: str
    """ISO 8601 timestamp of when this version was registered."""

    delta_size_mb: float
    """On-disk size of the LoRA adapter delta in megabytes."""

    active: bool = False
    """Whether this version is currently serving production traffic."""

    def to_dict(self) -> dict[str, object]:
        """Serialise to a JSON-compatible dict."""
        return {
            "version": self.version,
            "path": str(self.path),
            "accuracy": self.accuracy,
            "improvement_pp": self.improvement_pp,
            "n_training_samples": self.n_training_samples,
            "created_at": self.created_at,
            "delta_size_mb": self.delta_size_mb,
            "active": self.active,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> WeightManifest:
        """Deserialise from a dict loaded from JSON."""
        return cls(
            version=str(data["version"]),
            path=Path(str(data["path"])),
            accuracy=float(data["accuracy"]),
            improvement_pp=float(data.get("improvement_pp", 0.0)),
            n_training_samples=int(data["n_training_samples"]),
            created_at=str(data["created_at"]),
            delta_size_mb=float(data.get("delta_size_mb", 0.0)),
            active=bool(data.get("active", False)),
        )


class WeightRegistry:
    """Registry for ModernBERT LoRA adapter versions.

    Persists a JSON manifest and provides CRUD-like operations on registered
    adapter versions.

    Args:
        registry_dir: Directory where the manifest and adapter copies are stored.
            Defaults to ``~/.llmhosts/weights/modernbert/``.
    """

    def __init__(self, registry_dir: Path | None = None) -> None:
        self._dir = registry_dir or _DEFAULT_REGISTRY_DIR
        self._manifest_path = self._dir / _MANIFEST_FILENAME

    # ------------------------------------------------------------------
    # Core operations
    # ------------------------------------------------------------------

    def register(self, result: TrainingResult, version: str | None = None) -> WeightManifest:
        """Register a new adapter version in the manifest.

        Args:
            result: ``TrainingResult`` from ``LoRATrainer.train()``.
            version: Optional explicit version string.  Auto-generated from
                date and sample count when not supplied.

        Returns:
            The newly created ``WeightManifest`` entry.
        """
        self._dir.mkdir(parents=True, exist_ok=True)

        if version is None:
            version = _make_version(result.n_samples)

        manifest = WeightManifest(
            version=version,
            path=result.weights_path,
            accuracy=result.accuracy,
            improvement_pp=0.0,  # caller may update after evaluation
            n_training_samples=result.n_samples,
            created_at=datetime.datetime.now(datetime.timezone.utc).isoformat(),
            delta_size_mb=result.delta_size_mb,
            active=False,
        )

        versions = self._load_versions()
        # Replace existing entry with same version tag if present
        versions = [v for v in versions if v.version != version]
        versions.append(manifest)
        self._save_versions(versions)

        logger.info("Registered adapter version %s at %s", version, result.weights_path)
        return manifest

    def list_versions(self) -> list[WeightManifest]:
        """Return all registered adapter versions, newest first."""
        return sorted(self._load_versions(), key=lambda m: m.created_at, reverse=True)

    def get_active(self) -> WeightManifest | None:
        """Return the currently active adapter manifest, or ``None``."""
        for m in self._load_versions():
            if m.active:
                return m
        return None

    def activate(self, version: str) -> None:
        """Set the given version as active, deactivating all others.

        Args:
            version: Version string to activate (e.g. ``v20260224-10k``).

        Raises:
            ValueError: When the version is not found in the manifest.
        """
        versions = self._load_versions()
        found = False
        updated: list[WeightManifest] = []
        for m in versions:
            m.active = m.version == version
            if m.active:
                found = True
            updated.append(m)

        if not found:
            available = [m.version for m in versions]
            raise ValueError(f"Version '{version}' not found. Available: {available}")

        self._save_versions(updated)
        logger.info("Activated adapter version %s", version)

    # ------------------------------------------------------------------
    # Persistence
    # ------------------------------------------------------------------

    @property
    def registry_dir(self) -> Path:
        """Return the registry storage directory."""
        return self._dir

    def deactivate_all(self) -> None:
        """Deactivate all registered versions."""
        versions = self._load_versions()
        for v in versions:
            v.active = False
        self._save_versions(versions)

    def _load_versions(self) -> list[WeightManifest]:
        """Load the manifest from disk.  Returns empty list when not present."""
        if not self._manifest_path.exists():
            return []
        try:
            data = json.loads(self._manifest_path.read_text())
            return [WeightManifest.from_dict(entry) for entry in data]
        except Exception as exc:
            logger.warning("Failed to read manifest at %s: %s", self._manifest_path, exc)
            return []

    def _save_versions(self, versions: list[WeightManifest]) -> None:
        """Persist the manifest to disk."""
        self._dir.mkdir(parents=True, exist_ok=True)
        payload = [m.to_dict() for m in versions]
        self._manifest_path.write_text(json.dumps(payload, indent=2))


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _make_version(n_samples: int) -> str:
    """Generate a version string: ``v{YYYYMMDD}-{n_samples_k}k``.

    Examples:
        >>> _make_version(10_000)
        'v20260224-10k'
        >>> _make_version(500)
        'v20260224-0k'
    """
    date_str = datetime.date.today().strftime("%Y%m%d")
    k_samples = n_samples // 1000
    return f"v{date_str}-{k_samples}k"


def _is_valid_version(version: str) -> bool:
    """Return True if version matches the expected format ``v{YYYYMMDD}-{N}k``."""
    return bool(re.fullmatch(r"v\d{8}-\d+k", version))
