# bclm/hf_loader.py
"""
Model loading from Hugging Face, local directories, and URLs.

All weight loading uses safetensors for safety. The architecture
registry is consulted to instantiate the correct model class based
on the ``"architecture"`` field in ``config.json``.
"""
from __future__ import annotations

import json
import os
import sys
import tempfile
from pathlib import Path
from typing import Any, Dict, Optional, Tuple, Union
from urllib.parse import urlparse

import torch
import torch.nn as nn

from .config import (
    CONFIG_FILENAME,
    DEFAULT_HF_ORG,
    WEIGHTS_FILENAME,
)
from .registry import get_architecture


# =============================================================================
# Source detection
# =============================================================================

def _is_url(source: str) -> bool:
    parsed = urlparse(source)
    return parsed.scheme in ("http", "https")


def _is_local(source: str) -> bool:
    return os.path.exists(source)


def _is_hf_repo(source: str) -> bool:
    return not _is_url(source) and not _is_local(source)


def _normalise_hf_repo(source: str) -> str:
    """
    If *source* has no ``/``, prepend the default HF org.

    ``"bclm-1-small-preview"`` → ``"bclm/bclm-1-small-preview"``
    ``"myorg/mymodel"``        → ``"myorg/mymodel"``  (unchanged)
    """
    if "/" not in source:
        return f"{DEFAULT_HF_ORG}/{source}"
    return source


# =============================================================================
# File fetchers
# =============================================================================

def _fetch_from_local(
    directory: Union[str, Path],
) -> Tuple[Path, Path]:
    """Return (config_path, weights_path) from a local directory."""
    d = Path(directory)
    config_path = d / CONFIG_FILENAME
    weights_path = d / WEIGHTS_FILENAME

    if not config_path.is_file():
        raise FileNotFoundError(
            f"Config not found: {config_path}\n"
            f"Expected '{CONFIG_FILENAME}' in {d}"
        )
    if not weights_path.is_file():
        raise FileNotFoundError(
            f"Weights not found: {weights_path}\n"
            f"Expected '{WEIGHTS_FILENAME}' in {d}"
        )

    return config_path, weights_path


def _fetch_from_hf(
    repo_id: str,
    cache_dir: Optional[str] = None,
) -> Tuple[Path, Path]:
    """Download config + weights from Hugging Face Hub."""
    try:
        from huggingface_hub import hf_hub_download
    except ImportError as exc:
        raise ImportError(
            "Loading from Hugging Face requires 'huggingface_hub'.\n"
            "Install with:  pip install huggingface_hub"
        ) from exc

    repo_id = _normalise_hf_repo(repo_id)

    kwargs: Dict[str, Any] = {"repo_id": repo_id}
    if cache_dir:
        kwargs["cache_dir"] = cache_dir

    config_path = Path(hf_hub_download(filename=CONFIG_FILENAME, **kwargs))
    weights_path = Path(hf_hub_download(filename=WEIGHTS_FILENAME, **kwargs))

    return config_path, weights_path


def _download_file(url: str, dest: Path) -> None:
    """Download a single file via urllib (no extra dependencies)."""
    from urllib.request import urlopen, Request
    from urllib.error import HTTPError

    request = Request(url, headers={"User-Agent": "bclm-pytorch"})
    try:
        with urlopen(request, timeout=120) as resp, open(dest, "wb") as f:
            while True:
                chunk = resp.read(1 << 20)  # 1 MiB
                if not chunk:
                    break
                f.write(chunk)
    except HTTPError as exc:
        raise FileNotFoundError(
            f"Failed to download {url} — HTTP {exc.code}"
        ) from exc


def _fetch_from_url(
    base_url: str,
    cache_dir: Optional[str] = None,
) -> Tuple[Path, Path]:
    """Download config + weights from an HTTPS URL directory."""
    base = base_url.rstrip("/")
    config_url = f"{base}/{CONFIG_FILENAME}"
    weights_url = f"{base}/{WEIGHTS_FILENAME}"

    if cache_dir:
        download_dir = Path(cache_dir)
        download_dir.mkdir(parents=True, exist_ok=True)
    else:
        download_dir = Path(tempfile.mkdtemp(prefix="bclm_"))

    config_path = download_dir / CONFIG_FILENAME
    weights_path = download_dir / WEIGHTS_FILENAME

    if not config_path.is_file():
        _download_file(config_url, config_path)
    if not weights_path.is_file():
        _download_file(weights_url, weights_path)

    return config_path, weights_path


# =============================================================================
# Config parsing
# =============================================================================

def _load_config(config_path: Path) -> Tuple[str, dict]:
    """
    Read ``config.json`` and return ``(architecture_name, model_dict)``.

    The JSON must contain an ``"architecture"`` key and a ``"model"``
    key with the hyperparameters.
    """
    with open(config_path, "r", encoding="utf-8") as f:
        raw = json.load(f)

    if not isinstance(raw, dict):
        raise ValueError(
            f"Expected a JSON object in {config_path}, got {type(raw).__name__}"
        )

    architecture = raw.get("architecture")
    if not architecture:
        raise ValueError(
            f"Missing 'architecture' key in {config_path}"
        )

    model_dict = raw.get("model")
    if not isinstance(model_dict, dict):
        raise ValueError(
            f"Missing or invalid 'model' key in {config_path}"
        )

    return str(architecture), model_dict


# =============================================================================
# Weight loading
# =============================================================================

def _load_weights(
    weights_path: Path,
    device: torch.device,
) -> Dict[str, torch.Tensor]:
    """Load weights from a safetensors file."""
    try:
        from safetensors.torch import load_file
    except ImportError as exc:
        raise ImportError(
            "Loading model weights requires 'safetensors'.\n"
            "Install with:  pip install safetensors"
        ) from exc

    return load_file(str(weights_path), device=str(device))


# =============================================================================
# dtype resolution
# =============================================================================

def _resolve_dtype(
    dtype: Optional[torch.dtype],
    device: torch.device,
) -> torch.dtype:
    """
    Determine the working dtype.

    Default: bfloat16 on CUDA (matching training precision), float32
    on CPU (bfloat16 math is slow on most CPUs).
    """
    if dtype is not None:
        return dtype
    if device.type == "cuda":
        return torch.bfloat16
    return torch.float32


def _resolve_device(device: Optional[Union[str, torch.device]]) -> torch.device:
    """Resolve device string / None to a ``torch.device``."""
    if device is not None:
        return torch.device(device)
    if torch.cuda.is_available():
        return torch.device("cuda")
    return torch.device("cpu")


# =============================================================================
# Main loader
# =============================================================================

def load_model(
    source: str,
    *,
    device: Optional[Union[str, torch.device]] = None,
    dtype: Optional[torch.dtype] = None,
    compile: bool = False,
    cache_dir: Optional[str] = None,
) -> Dict[str, Any]:
    """
    Load a BCLM model from any supported source.

    Parameters
    ----------
    source : str
        One of:
        - A short model name (e.g. ``"bclm-1-small-preview"``)
        - A full HF repo ID (e.g. ``"bclm/bclm-1-small-preview"``)
        - A local directory path containing config.json + model.safetensors
        - An HTTPS URL to a directory with those files
    device : str or torch.device, optional
        Target device. Auto-selects CUDA if available.
    dtype : torch.dtype, optional
        Weight dtype. Defaults to bfloat16 on CUDA, float32 on CPU.
    compile : bool
        Wrap the model with ``torch.compile`` for faster inference.
        Disabled by default to avoid warmup latency.
    cache_dir : str, optional
        Cache directory for downloaded files (HF hub or URL).

    Returns
    -------
    dict
        Keys: ``"model"``, ``"config"``, ``"generator"``,
        ``"architecture"``, ``"tokenizer_spec"``, ``"device"``,
        ``"dtype"``.
    """
    resolved_device = _resolve_device(device)
    resolved_dtype = _resolve_dtype(dtype, resolved_device)

    # ── Fetch files ─────────────────────────────────────────

    if _is_local(source):
        config_path, weights_path = _fetch_from_local(source)
        source_label = str(source)
    elif _is_url(source):
        config_path, weights_path = _fetch_from_url(source, cache_dir)
        source_label = source
    elif _is_hf_repo(source):
        config_path, weights_path = _fetch_from_hf(source, cache_dir)
        source_label = _normalise_hf_repo(source)
    else:
        raise ValueError(f"Cannot determine source type for: {source!r}")

    # ── Parse config ────────────────────────────────────────

    architecture_name, model_dict = _load_config(config_path)
    tokenizer_spec = model_dict.get("tokenizer", "")

    # ── Resolve architecture from registry ──────────────────

    entry = get_architecture(architecture_name)
    model_config = entry.config_class.from_dict(model_dict)

    # ── Instantiate model ───────────────────────────────────

    # Build on CPU first to avoid fragmented GPU allocation
    raw_model = entry.model_class(model_config)

    # ── Load weights ────────────────────────────────────────

    state_dict = _load_weights(weights_path, device=torch.device("cpu"))
    raw_model.load_state_dict(state_dict, strict=True)

    # Free the state dict before moving to device
    del state_dict

    # ── dtype + device ──────────────────────────────────────

    raw_model = raw_model.to(device=resolved_device, dtype=resolved_dtype)
    raw_model.eval()

    # ── Optional compilation ────────────────────────────────

    if compile:
        try:
            raw_model = torch.compile(raw_model)
        except Exception as exc:
            print(
                f"[bclm] torch.compile failed, continuing without: {exc}",
                file=sys.stderr,
            )

    # ── Wrap in inference class ─────────────────────────────

    generator = entry.inference_class(raw_model)

    param_count = sum(p.numel() for p in raw_model.parameters())
    print(
        f"[bclm] Loaded {architecture_name} from {source_label}\n"
        f"[bclm] Parameters: {param_count / 1e6:.2f}M  |  "
        f"Device: {resolved_device}  |  Dtype: {resolved_dtype}",
        flush=True,
    )

    return {
        "model": raw_model,
        "config": model_config,
        "generator": generator,
        "architecture": architecture_name,
        "tokenizer_spec": tokenizer_spec,
        "device": resolved_device,
        "dtype": resolved_dtype,
    }