"""
scrylo/core/model_manager.py
-----------------------------
Handles Phi-3 Mini Q4 download, caching, and verification.
Model is NEVER bundled — downloaded once to ~/.scrylo/models/ on first
`scrylo pitch` call and cached permanently. All future runs skip download.

Memory contract (enforced here and in llm_adapter.py):
  - NLP models must be unloaded BEFORE this module loads the LLM
  - LLM must be deleted (del llm) AFTER batch generation completes
  - Never load LLM and NLP simultaneously on 4 GB RAM
"""

from __future__ import annotations

import hashlib
import shutil
from pathlib import Path
from typing import Callable

from rich.console import Console
from rich.progress import (
    BarColumn,
    DownloadColumn,
    Progress,
    TextColumn,
    TimeRemainingColumn,
    TransferSpeedColumn,
)

from scrylo.core.config import MODELS_DIR, settings

console = Console()

# ── Model registry ────────────────────────────────────────────────────────
# Add future models here — llm_adapter.py picks up via config["llm"]["model"]
MODEL_REGISTRY: dict[str, dict] = {
    "phi3-mini-q4_k_m.gguf": {
        "repo_id": "microsoft/Phi-3-mini-4k-instruct-gguf",
        "filename": "Phi-3-mini-4k-instruct-q4.gguf",
        "size_gb": 2.2,
        "description": "Phi-3 Mini Q4_K_M — default, runs on 4 GB RAM",
        # SHA256 populated after first successful download for future verification
        "sha256": None,
    },
    "mistral-7b-q4.gguf": {
        "repo_id": "TheBloke/Mistral-7B-Instruct-v0.2-GGUF",
        "filename": "mistral-7b-instruct-v0.2.Q4_K_M.gguf",
        "size_gb": 4.1,
        "description": "Mistral 7B Q4 — only for 8 GB+ RAM systems",
        "sha256": None,
    },
}


def model_path(model_name: str | None = None) -> Path:
    """Return the full path for a model file."""
    name = model_name or settings["llm"]["model"]
    return MODELS_DIR / name


def is_downloaded(model_name: str | None = None) -> bool:
    """Check if the model file exists and is non-empty."""
    p = model_path(model_name)
    return p.exists() and p.stat().st_size > 100_000_000  # > 100 MB sanity check


def get_sha256(path: Path, progress_cb: Callable[[int], None] | None = None) -> str:
    """Compute SHA256 of a file. Optionally report bytes processed."""
    h = hashlib.sha256()
    chunk = 1024 * 1024  # 1 MB
    with path.open("rb") as f:
        while True:
            block = f.read(chunk)
            if not block:
                break
            h.update(block)
            if progress_cb:
                progress_cb(len(block))
    return h.hexdigest()


def ensure_model(
    model_name: str | None = None,
    auto_confirm: bool = False,
) -> Path:
    """
    Main entry point — called by llm_adapter before loading the model.

    1. Check if model is already cached → return path immediately.
    2. If not, show a one-time setup prompt with size warning.
    3. Download with a live Rich progress bar.
    4. Compute and store SHA256 for future integrity checks.
    5. If user declines → raise ModelNotAvailableError (caller activates Jinja2 fallback).
    """
    name = model_name or settings["llm"]["model"]
    dest = model_path(name)

    if is_downloaded(name):
        return dest

    info = MODEL_REGISTRY.get(name, {})
    size = info.get("size_gb", "?")
    desc = info.get("description", name)

    console.print()
    console.print(
        f"[bold yellow]⚡ First-time AI setup required[/bold yellow]\n"
        f"   Scrylo needs to download the AI model to generate personalised pitches.\n\n"
        f"   [bold]{desc}[/bold]\n"
        f"   Size: [cyan]{size} GB[/cyan] — downloaded once, cached permanently at:\n"
        f"   [dim]{dest}[/dim]\n"
    )

    if not auto_confirm:
        confirm = console.input("   Download now? [Y/n] ").strip().lower()
        if confirm not in ("", "y", "yes"):
            raise ModelNotAvailableError(
                "Model download declined. Run with --no-llm to use template fallback."
            )

    _download(name, dest, info)
    return dest


def _download(name: str, dest: Path, info: dict) -> None:
    """Download model from HuggingFace Hub with a live progress bar."""
    try:
        from huggingface_hub import hf_hub_download
    except ImportError:
        raise RuntimeError(
            "huggingface_hub not installed. Run: pip install huggingface-hub"
        )

    repo_id = info.get("repo_id", settings["llm"]["repo_id"])
    filename = info.get("filename", name)

    console.print(f"\n[dim]Downloading from HuggingFace: {repo_id}/{filename}[/dim]")

    with Progress(
        TextColumn("[bold green]{task.description}"),
        BarColumn(),
        DownloadColumn(),
        TransferSpeedColumn(),
        TimeRemainingColumn(),
        console=console,
    ) as progress:
        task = progress.add_task(f"Downloading {name}", total=None)

        tmp = hf_hub_download(
            repo_id=repo_id,
            filename=filename,
            local_dir=MODELS_DIR,
            local_dir_use_symlinks=False,
        )
        progress.update(task, completed=100, total=100)

    # hf_hub_download may save under original filename — rename to our key
    tmp_path = Path(tmp)
    if tmp_path != dest:
        shutil.move(str(tmp_path), dest)

    console.print(f"[bold green]✓ Model saved to {dest}[/bold green]")
    console.print("[dim]Computing checksum…[/dim]")
    sha = get_sha256(dest)
    console.print(f"[dim]SHA256: {sha}[/dim]")

    # Persist SHA256 for future integrity checks
    checksum_file = dest.with_suffix(".sha256")
    checksum_file.write_text(sha)
    console.print("[bold green]✓ Setup complete — model ready for all future runs[/bold green]\n")


def verify_integrity(model_name: str | None = None) -> bool:
    """
    Check model file against stored SHA256.
    Returns True if OK, False if corrupt/missing.
    Called automatically before loading — corrupt downloads are re-fetched.
    """
    name = model_name or settings["llm"]["model"]
    dest = model_path(name)
    checksum_file = dest.with_suffix(".sha256")

    if not dest.exists():
        return False
    if not checksum_file.exists():
        # No checksum stored yet — compute and store it now
        sha = get_sha256(dest)
        checksum_file.write_text(sha)
        return True

    stored = checksum_file.read_text().strip()
    actual = get_sha256(dest)
    return stored == actual


def delete_model(model_name: str | None = None) -> None:
    """Remove a cached model to free disk space."""
    name = model_name or settings["llm"]["model"]
    dest = model_path(name)
    if dest.exists():
        dest.unlink()
        console.print(f"[yellow]Deleted model: {dest}[/yellow]")
    checksum = dest.with_suffix(".sha256")
    if checksum.exists():
        checksum.unlink()


def list_models() -> list[dict]:
    """Return status of all known models — for CLI `scrylo models` command."""
    result = []
    for name, info in MODEL_REGISTRY.items():
        dest = model_path(name)
        downloaded = is_downloaded(name)
        size_on_disk = round(dest.stat().st_size / 1e9, 2) if downloaded else 0
        result.append({
            "name": name,
            "description": info["description"],
            "size_gb": info["size_gb"],
            "downloaded": downloaded,
            "size_on_disk_gb": size_on_disk,
            "path": str(dest) if downloaded else None,
        })
    return result


class ModelNotAvailableError(Exception):
    """
    Raised when model download is declined or fails.
    Callers catch this and activate the Jinja2 template fallback.
    """
    pass
