"""llama-cpp-python local AI provider — runs GGUF models fully in-process.

NO network connections at all — no TCP, no localhost, no ports.
This is the fallback when Ollama is blocked by CSCA exam software
(which blocks all network ports including localhost:11434).

Supports:
  - Text generation via llama_cpp.Llama
  - Vision (image+text) via llama_cpp.Llama with clip_model_path (MiniCPM-V)

Model files are stored in ~/.muravey-china/models/:
  - Vision: ggml-model-Q4_K_M.gguf + mmproj-model-f16.gguf
  - Text:   qwen2.5-1.5b-instruct-q4_k_m.gguf

Usage:
    provider = LlamaCppProvider()
    text = provider.generate("What is 2+2?")
    text = provider.generate("Describe this image", image_base64="...")
"""

from __future__ import annotations

import base64
import logging
import os
import threading
from pathlib import Path

log = logging.getLogger(__name__)

# Model directory
MODELS_DIR = Path(os.path.expanduser("~/.muravey-china/models"))

# Expected model filenames (local names after download)
VISION_MODEL_FILE = "ggml-model-Q4_K_M.gguf"
VISION_MMPROJ_FILE = "mmproj-model-f16.gguf"

# Text models — priority order: Qwen 14B (best) → Qwen 7B → Qwen 1.5B (fallback)
TEXT_MODEL_FILE_14B = "Qwen2.5-14B-Instruct-Q3_K_S.gguf"  # ~6.7GB, best for Chinese exams
TEXT_MODEL_FILE = "qwen2.5-7b-instruct-q3_k_m.gguf"     # ~3.8GB, good quality
TEXT_MODEL_FILE_SMALL = "qwen2.5-1.5b-instruct-q4_k_m.gguf"  # ~1.1GB, fallback

# HuggingFace download URLs
_HF_REPO_VISION = "openbmb/MiniCPM-V-2_6-gguf"
_HF_REPO_TEXT_7B = "Qwen/Qwen2.5-7B-Instruct-GGUF"
_HF_REPO_TEXT_SMALL = "Qwen/Qwen2.5-1.5B-Instruct-GGUF"

GGUF_DOWNLOAD_URLS: dict[str, str] = {
    VISION_MODEL_FILE: f"https://huggingface.co/{_HF_REPO_VISION}/resolve/main/{VISION_MODEL_FILE}",
    VISION_MMPROJ_FILE: f"https://huggingface.co/{_HF_REPO_VISION}/resolve/main/{VISION_MMPROJ_FILE}",
    TEXT_MODEL_FILE: f"https://huggingface.co/{_HF_REPO_TEXT_7B}/resolve/main/{TEXT_MODEL_FILE}",
}


def is_llamacpp_available() -> bool:
    """Check if llama-cpp-python is importable."""
    try:
        import llama_cpp  # noqa: F401
        return True
    except Exception:
        return False


def _find_text_model() -> tuple[Path, str] | None:
    """Return (path, chat_format) for best available text model.

    Priority: Qwen 14B (best) → Qwen 7B → Qwen 1.5B (fallback).
    """
    # 1. Qwen 14B — best quality for Chinese exams
    p14b = MODELS_DIR / TEXT_MODEL_FILE_14B
    if p14b.is_file():
        return p14b, "chatml"
    # 2. Qwen 7B
    p7b = MODELS_DIR / TEXT_MODEL_FILE
    if p7b.is_file():
        return p7b, "chatml"
    # 3. Qwen 1.5B fallback
    p_small = MODELS_DIR / TEXT_MODEL_FILE_SMALL
    if p_small.is_file():
        return p_small, "chatml"
    return None


def are_models_present(vision: bool = True, text: bool = True) -> bool:
    """Check if required GGUF model files exist on disk."""
    if vision:
        if not (MODELS_DIR / VISION_MODEL_FILE).is_file():
            return False
        if not (MODELS_DIR / VISION_MMPROJ_FILE).is_file():
            return False
    if text:
        if _find_text_model() is None:  # returns tuple or None
            return False
    return True


class LlamaCppProvider:
    """In-process GGUF model provider using llama-cpp-python.

    Loads models into memory once. All inference is done in-process
    with ZERO network connections — safe under CSCA port blocking.
    """

    def __init__(self) -> None:
        self._vision_model = None  # llama_cpp.Llama instance for vision
        self._text_model = None    # llama_cpp.Llama instance for text
        self._vision_lock = threading.Lock()
        self._text_lock = threading.Lock()
        self._loaded_vision = False
        self._loaded_text = False

    # ------------------------------------------------------------------ #
    # Model loading
    # ------------------------------------------------------------------ #

    def load_vision_model(self) -> bool:
        """Load the vision GGUF model + mmproj into memory.

        Returns True if loaded successfully, False otherwise.
        """
        if self._loaded_vision:
            return True

        with self._vision_lock:
            if self._loaded_vision:
                return True

            model_path = MODELS_DIR / VISION_MODEL_FILE
            mmproj_path = MODELS_DIR / VISION_MMPROJ_FILE

            if not model_path.is_file() or not mmproj_path.is_file():
                log.warning(
                    "[llamacpp] Vision model files not found in %s", MODELS_DIR
                )
                return False

            try:
                from llama_cpp import Llama
                from llama_cpp.llama_chat_format import MiniCPMv26ChatHandler

                chat_handler = MiniCPMv26ChatHandler(
                    clip_model_path=str(mmproj_path),
                )

                self._vision_model = Llama(
                    model_path=str(model_path),
                    chat_handler=chat_handler,
                    n_ctx=2048,
                    n_gpu_layers=0,  # CPU only — safe under CSCA
                    verbose=False,
                )
                self._loaded_vision = True
                log.info(
                    "[llamacpp] Vision model loaded: %s (GPU offload enabled)",
                    model_path.name,
                )
                return True
            except Exception as exc:
                log.error("[llamacpp] Failed to load vision model: %s", exc)
                # Try fallback without specific chat handler
                try:
                    from llama_cpp import Llama

                    self._vision_model = Llama(
                        model_path=str(model_path),
                        chat_format="chatml",
                        clip_model_path=str(mmproj_path),
                        n_ctx=2048,
                        n_gpu_layers=0,
                        verbose=False,
                    )
                    self._loaded_vision = True
                    log.info(
                        "[llamacpp] Vision model loaded with chatml fallback: %s",
                        model_path.name,
                    )
                    return True
                except Exception as exc2:
                    log.error("[llamacpp] Vision model fallback also failed: %s", exc2)
                    return False

    def load_text_model(self) -> bool:
        """Load the text GGUF model into memory.

        Returns True if loaded successfully, False otherwise.
        """
        if self._loaded_text:
            return True

        with self._text_lock:
            if self._loaded_text:
                return True

            found = _find_text_model()

            if found is None:
                log.warning(
                    "[llamacpp] No text model file found in %s", MODELS_DIR
                )
                return False

            model_path, chat_format = found

            try:
                from llama_cpp import Llama

                self._text_model = Llama(
                    model_path=str(model_path),
                    chat_format=chat_format,
                    n_ctx=2048,       # Enough for exam questions
                    n_batch=512,      # Faster prompt processing
                    n_gpu_layers=0,   # CPU only — safe under CSCA
                    verbose=False,
                )
                self._loaded_text = True
                log.info(
                    "[llamacpp] Text model loaded: %s (format=%s, GPU offload enabled)",
                    model_path.name, chat_format,
                )
                return True
            except Exception as exc:
                log.error("[llamacpp] Failed to load text model: %s", exc)
                return False

    def load_all(self) -> bool:
        """Load both vision and text models. Returns True if at least one loaded."""
        v = self.load_vision_model()
        t = self.load_text_model()
        return v or t

    # ------------------------------------------------------------------ #
    # Inference
    # ------------------------------------------------------------------ #

    def generate(
        self,
        prompt: str,
        image_base64: str | None = None,
        *,
        system_prompt: str = "You are a helpful assistant.",
        max_tokens: int = 512,
        temperature: float = 0.0,
    ) -> str:
        """Generate text from prompt, optionally with an image.

        This method does NOT use any network connections.

        Args:
            prompt: The user prompt text.
            image_base64: Optional base64-encoded image (JPEG/PNG).
                          If provided, uses the vision model.
            system_prompt: System instruction.
            max_tokens: Maximum tokens to generate.
            temperature: Sampling temperature.

        Returns:
            Generated text string, or empty string on failure.
        """
        if image_base64:
            return self._generate_vision(
                prompt, image_base64,
                system_prompt=system_prompt,
                max_tokens=max_tokens,
                temperature=temperature,
            )
        else:
            return self._generate_text(
                prompt,
                system_prompt=system_prompt,
                max_tokens=max_tokens,
                temperature=temperature,
            )

    def _generate_vision(
        self,
        prompt: str,
        image_base64: str,
        *,
        system_prompt: str,
        max_tokens: int,
        temperature: float,
    ) -> str:
        """Run vision inference with image + text prompt."""
        if not self._loaded_vision:
            if not self.load_vision_model():
                log.error("[llamacpp] Vision model not available")
                return ""

        # Build data URI for the image
        data_uri = f"data:image/jpeg;base64,{image_base64}"

        messages = [
            {"role": "system", "content": system_prompt},
            {
                "role": "user",
                "content": [
                    {"type": "image_url", "image_url": {"url": data_uri}},
                    {"type": "text", "text": prompt},
                ],
            },
        ]

        try:
            result = self._vision_model.create_chat_completion(
                messages=messages,
                max_tokens=max_tokens,
                temperature=temperature,
            )
            text = result["choices"][0]["message"]["content"] or ""
            log.info("[llamacpp] Vision response: %d chars", len(text))
            return text.strip()
        except Exception as exc:
            log.error("[llamacpp] Vision inference failed: %s", exc)
            return ""

    def _generate_text(
        self,
        prompt: str,
        *,
        system_prompt: str,
        max_tokens: int,
        temperature: float,
    ) -> str:
        """Run text-only inference."""
        if not self._loaded_text:
            if not self.load_text_model():
                log.error("[llamacpp] Text model not available")
                return ""

        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": prompt},
        ]

        try:
            result = self._text_model.create_chat_completion(
                messages=messages,
                max_tokens=max_tokens,
                temperature=temperature,
            )
            text = result["choices"][0]["message"]["content"] or ""
            log.info("[llamacpp] Text response: %d chars", len(text))
            return text.strip()
        except Exception as exc:
            log.error("[llamacpp] Text inference failed: %s", exc)
            return ""

    # ------------------------------------------------------------------ #
    # Status
    # ------------------------------------------------------------------ #

    @property
    def vision_ready(self) -> bool:
        return self._loaded_vision

    @property
    def text_ready(self) -> bool:
        return self._loaded_text

    def __repr__(self) -> str:
        return (
            f"<LlamaCppProvider vision={self._loaded_vision} text={self._loaded_text}>"
        )


# ===================================================================== #
# Singleton instance — created lazily, shared across the app
# ===================================================================== #

_instance: LlamaCppProvider | None = None
_instance_lock = threading.Lock()


def get_llamacpp_provider() -> LlamaCppProvider | None:
    """Return the global LlamaCppProvider instance, or None if unavailable.

    Creates and loads models on first call. Thread-safe.
    """
    global _instance
    if _instance is not None:
        return _instance

    if not is_llamacpp_available():
        log.debug("[llamacpp] llama-cpp-python not installed — provider unavailable")
        return None

    if not are_models_present(vision=True, text=True) and not are_models_present(vision=True, text=False):
        log.debug("[llamacpp] No GGUF model files found in %s", MODELS_DIR)
        return None

    with _instance_lock:
        if _instance is not None:
            return _instance
        _instance = LlamaCppProvider()
        _instance.load_all()

    return _instance


def download_gguf_models(progress_callback=None) -> bool:
    """Download GGUF model files from HuggingFace to MODELS_DIR.

    This DOES require network — call this BEFORE the exam, not during.

    Args:
        progress_callback: Optional callable(filename, bytes_done, bytes_total).

    Returns:
        True if all files downloaded successfully.
    """
    import urllib.request

    MODELS_DIR.mkdir(parents=True, exist_ok=True)
    success = True

    total_files = len(GGUF_DOWNLOAD_URLS)
    print(f"Download directory: {MODELS_DIR}")
    print(f"Files to check: {total_files}")
    print()

    for i, (filename, url) in enumerate(GGUF_DOWNLOAD_URLS.items(), 1):
        dest = MODELS_DIR / filename
        if dest.is_file():
            size_mb = dest.stat().st_size / (1024 * 1024)
            print(f"[{i}/{total_files}] {filename} — already exists ({size_mb:.0f} MB)")
            continue

        print(f"[{i}/{total_files}] Downloading {filename}...")
        print(f"  URL: {url}")
        try:
            req = urllib.request.Request(url, headers={"User-Agent": "MuraveyChina/1.0"})
            with urllib.request.urlopen(req, timeout=600) as resp:
                total = int(resp.headers.get("Content-Length", 0))
                total_mb = total / (1024 * 1024) if total else 0
                done = 0
                tmp_path = dest.with_suffix(".tmp")
                with open(tmp_path, "wb") as f:
                    while True:
                        chunk = resp.read(1024 * 1024)  # 1MB chunks
                        if not chunk:
                            break
                        f.write(chunk)
                        done += len(chunk)
                        done_mb = done / (1024 * 1024)
                        if total:
                            pct = done * 100 // total
                            print(f"\r  Progress: {done_mb:.0f}/{total_mb:.0f} MB ({pct}%)", end="", flush=True)
                        else:
                            print(f"\r  Downloaded: {done_mb:.0f} MB", end="", flush=True)
                        if progress_callback:
                            progress_callback(filename, done, total)
                print()  # newline after progress
                # Verify file size matches expected
                if total and done < total:
                    print(f"  WARNING: Incomplete download ({done}/{total} bytes) — deleting")
                    tmp_path.unlink(missing_ok=True)
                    success = False
                    continue
                tmp_path.rename(dest)
                print(f"  Done: {filename} ({done / (1024*1024):.0f} MB)")
        except Exception as exc:
            print(f"\n  ERROR: {exc}")
            log.error("[llamacpp] Failed to download %s: %s", filename, exc)
            success = False

    print()
    if success:
        print("All models ready!")
    else:
        print("Some downloads failed. Retry with internet connection.")
    return success
