"""Automatic weight update checker for ModernBERT routing adapters.

``WeightUpdater`` checks a remote manifest (HuggingFace Hub by default) for
newer fine-tuned adapter versions and downloads them if found.

All HTTP errors are handled silently — the updater should never disrupt the
serving path if the update check fails.
"""

from __future__ import annotations

import dataclasses
import json
import logging
from pathlib import Path
from typing import Any

import httpx

from llmhosts.training.registry import WeightRegistry

logger = logging.getLogger(__name__)

_DEFAULT_REGISTRY_URL: str = "https://huggingface.co/llmhosts/modernbert-router/resolve/main/manifest.json"


@dataclasses.dataclass
class UpdateResult:
    """Result of an update check."""

    updated: bool
    """True when new weights were downloaded and registered successfully."""

    version: str | None
    """Version string of the newly downloaded adapter (when ``updated=True``)."""

    message: str
    """Human-readable description of what happened."""


class WeightUpdater:
    """Checks for and downloads newer ModernBERT adapter versions.

    Compares the remote manifest (e.g. on HuggingFace Hub) with the locally
    active version.  If a newer version is available, downloads and registers
    the adapter weights using ``WeightRegistry``.

    All network errors are caught silently — the serving path is unaffected by
    update check failures.

    Args:
        registry: Optional ``WeightRegistry`` instance.  A default registry
            (``~/.llmhosts/weights/modernbert/``) is created when not supplied.
    """

    def __init__(self, registry: WeightRegistry | None = None) -> None:
        self._registry = registry or WeightRegistry()

    def check_and_update(
        self,
        registry_url: str | None = None,
        timeout_s: float = 10.0,
    ) -> UpdateResult:
        """Check for a newer adapter version and download if available.

        Steps:
        1. Fetch ``manifest.json`` from ``registry_url``.
        2. Parse the manifest and find the latest entry.
        3. Compare with the locally active version.
        4. If newer, download the adapter weights and register locally.

        Args:
            registry_url: URL of the remote ``manifest.json``.  Defaults to
                the official HuggingFace Hub location.
            timeout_s: HTTP request timeout in seconds.

        Returns:
            An ``UpdateResult`` describing what happened.
        """
        url = registry_url or _DEFAULT_REGISTRY_URL

        # -- Fetch remote manifest --
        try:
            response = httpx.get(url, timeout=timeout_s, follow_redirects=True)
            response.raise_for_status()
            remote_data = response.json()
        except httpx.TimeoutException:
            msg = f"Timeout ({timeout_s}s) fetching remote manifest from {url}"
            logger.debug(msg)
            return UpdateResult(updated=False, version=None, message=msg)
        except httpx.HTTPStatusError as exc:
            msg = f"HTTP {exc.response.status_code} fetching remote manifest from {url}"
            logger.debug(msg)
            return UpdateResult(updated=False, version=None, message=msg)
        except httpx.RequestError as exc:
            msg = f"Network error fetching remote manifest: {exc}"
            logger.debug(msg)
            return UpdateResult(updated=False, version=None, message=msg)
        except (json.JSONDecodeError, ValueError) as exc:
            msg = f"Malformed JSON in remote manifest: {exc}"
            logger.debug(msg)
            return UpdateResult(updated=False, version=None, message=msg)

        # -- Parse remote manifest --
        try:
            remote_versions = self._parse_remote_manifest(remote_data)
        except Exception as exc:
            msg = f"Could not parse remote manifest: {exc}"
            logger.debug(msg)
            return UpdateResult(updated=False, version=None, message=msg)

        if not remote_versions:
            return UpdateResult(updated=False, version=None, message="Remote manifest is empty.")

        # -- Find newest remote version --
        latest_remote = max(remote_versions, key=lambda v: str(v.get("created_at", "")))
        remote_version_str = str(latest_remote.get("version", ""))

        # -- Compare with local active version --
        local_active = self._registry.get_active()

        local_created = local_active.created_at if local_active else ""
        remote_created = str(latest_remote.get("created_at", ""))
        if local_created and local_created >= remote_created:
            return UpdateResult(updated=False, version=None, message="Already at latest version")

        # -- Download adapter weights --
        weights_url = str(latest_remote.get("weights_url", ""))
        if not weights_url:
            msg = "Remote manifest entry missing 'weights_url'; cannot download."
            logger.debug(msg)
            return UpdateResult(updated=False, version=None, message=msg)

        try:
            result = self._download_adapter(weights_url, remote_version_str, latest_remote, timeout_s)
            return result
        except Exception as exc:
            msg = f"Failed to download adapter weights: {exc}"
            logger.debug(msg)
            return UpdateResult(updated=False, version=None, message=msg)

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

    @staticmethod
    def _parse_remote_manifest(data: Any) -> list[dict[str, Any]]:
        """Parse the remote manifest JSON into a list of version dicts.

        Supports both a plain list and a dict with a ``versions`` key.
        """
        if isinstance(data, list):
            return [entry for entry in data if isinstance(entry, dict)]
        if isinstance(data, dict):
            versions = data.get("versions", data.get("entries", []))
            if isinstance(versions, list):
                return [e for e in versions if isinstance(e, dict)]
        raise ValueError(f"Unexpected manifest format: {type(data)}")

    def _download_adapter(
        self,
        weights_url: str,
        version: str,
        manifest_entry: dict[str, Any],
        timeout_s: float,
    ) -> UpdateResult:
        """Download adapter weights from ``weights_url`` and register locally."""
        registry_dir = self._registry.registry_dir
        adapter_dir = registry_dir / version
        adapter_dir.mkdir(parents=True, exist_ok=True)

        # Download as a tar/zip or single file
        weights_file = adapter_dir / "adapter_model.safetensors"

        logger.info("Downloading ModernBERT adapter %s from %s", version, weights_url)
        with httpx.stream("GET", weights_url, timeout=timeout_s, follow_redirects=True) as stream:
            stream.raise_for_status()
            with weights_file.open("wb") as fh:
                for chunk in stream.iter_bytes(chunk_size=65536):
                    fh.write(chunk)

        # Write minimal adapter config if not included in download
        config_file = adapter_dir / "adapter_config.json"
        if not config_file.exists():
            import json as _json

            config_file.write_text(
                _json.dumps(
                    {
                        "base_model_name_or_path": manifest_entry.get("base_model", "answerdotai/ModernBERT-base"),
                        "peft_type": "LORA",
                        "r": manifest_entry.get("lora_rank", 16),
                        "lora_alpha": manifest_entry.get("lora_alpha", 32),
                    },
                    indent=2,
                )
            )

        # Register and activate the new version
        from llmhosts.training.trainer import TrainingResult

        size_mb = weights_file.stat().st_size / (1024 * 1024) if weights_file.exists() else 0.0

        synthetic_result = TrainingResult(
            weights_path=adapter_dir,
            accuracy=float(manifest_entry.get("accuracy", 0.0)),
            latency_ms=float(manifest_entry.get("latency_ms", 0.0)),
            n_samples=int(manifest_entry.get("n_training_samples", 0)),
            epochs_completed=int(manifest_entry.get("epochs", 0)),
            delta_size_mb=round(size_mb, 2),
        )

        manifest = self._registry.register(synthetic_result, version=version)
        self._registry.activate(manifest.version)

        msg = f"Downloaded and activated adapter version {version}."
        logger.info(msg)
        return UpdateResult(updated=True, version=version, message=msg)


def _resolve_download_dir() -> Path:
    """Return the default directory for downloaded adapter weights."""
    return Path.home() / ".llmhosts" / "weights" / "modernbert"
