model_engine.py

"""
Nyra model engine.

- Local / non-Space: mock streaming only. NEVER downloads GGUF weights.
- Hugging Face Spaces: load GGUF from Hub cache/preload and stream via llama-cpp.
"""

from __future__ import annotations

import ctypes
import os
import re
import subprocess
import sys
import threading
import time
from pathlib import Path
from typing import Generator, Iterable, List, Optional

ROOT = Path(__file__).resolve().parent
PROMPTS_DIR = ROOT / "prompts"

REPO_ID = "HauhauCS/Qwen3.6-35B-A3B-Uncensored-HauhauCS-Aggressive"
# Prefer smaller quant first: faster cold-load + fits ZeroGPU time/quota better.
GGUF_FILENAME = (
    "Qwen3.6-35B-A3B-Uncensored-HauhauCS-Aggressive-IQ4_XS.gguf"
)
FALLBACK_FILENAMES = [
    "Qwen3.6-35B-A3B-Uncensored-HauhauCS-Aggressive-IQ3_M.gguf",
    "Qwen3.6-35B-A3B-Uncensored-HauhauCS-Aggressive-Q4_K_M.gguf",
    "Qwen3.6-35B-A3B-Uncensored-HauhauCS-Aggressive-Q4_K_P.gguf",
]

N_CTX = int(os.getenv("NYRA_N_CTX", "2048"))
N_GPU_LAYERS = int(os.getenv("NYRA_N_GPU_LAYERS", "-1"))
DEFAULT_MAX_TOKENS = int(os.getenv("NYRA_MAX_TOKENS", "512"))
DEFAULT_TEMPERATURE = float(os.getenv("NYRA_TEMPERATURE", "0.7"))

FORCE_LOCAL_MODEL = os.getenv("FORCE_LOCAL_MODEL", "").strip() in {
    "1",
    "true",
    "yes",
}

_llm = None
_llm_lock = threading.Lock()
_llm_error: Optional[str] = None
_llama_ready = False
_warmup_started = False


def is_spaces() -> bool:
    return bool(
        os.getenv("SPACE_ID")
        or os.getenv("SPACE_HOST")
        or os.getenv("SYSTEM") == "spaces"
    )


def should_load_model() -> bool:
    return is_spaces() or FORCE_LOCAL_MODEL


def runtime_mode() -> str:
    if is_spaces():
        return "spaces"
    if FORCE_LOCAL_MODEL:
        return "local-forced"
    return "mock"


def load_system_prompt(thinking: bool = False) -> str:
    name = "nyra_system_thinking.md" if thinking else "nyra_system.md"
    path = PROMPTS_DIR / name
    if path.exists():
        return path.read_text(encoding="utf-8").strip()
    return (
        "You are Nyra, a helpful chat assistant on a Hugging Face Space. "
        "You are not Grok and not affiliated with xAI. "
        "Respond in the user's language."
    )


def _preload_cuda_from_torch() -> Optional[Path]:
    """
    ZeroGPU / Spaces often lack system libcudart, but PyTorch ships CUDA libs.
    Preload them so llama-cpp CUDA wheels can resolve symbols.
    """
    try:
        import torch
    except ImportError:
        return None

    torch_lib = Path(torch.__file__).resolve().parent / "lib"
    if not torch_lib.is_dir():
        return None

    # Ensure dynamic linker search path
    current = os.environ.get("LD_LIBRARY_PATH", "")
    parts = [str(torch_lib)] + ([current] if current else [])
    os.environ["LD_LIBRARY_PATH"] = ":".join(parts)

    # Also common CUDA paths on Spaces images
    for extra in (
        "/usr/local/cuda/lib64",
        "/usr/local/cuda/lib",
        "/usr/lib/x86_64-linux-gnu",
    ):
        if Path(extra).is_dir():
            os.environ["LD_LIBRARY_PATH"] = (
                f"{extra}:{os.environ['LD_LIBRARY_PATH']}"
            )

    candidates = [
        "libcudart.so.12",
        "libcudart.so.11.0",
        "libcudart.so",
        "libcublas.so.12",
        "libcublas.so",
        "libnvrtc.so.12",
        "libnvrtc.so",
    ]
    for name in candidates:
        path = torch_lib / name
        if path.exists():
            try:
                ctypes.CDLL(str(path), mode=ctypes.RTLD_GLOBAL)
            except OSError:
                continue
    return torch_lib


def _find_gguf_path() -> Optional[Path]:
    env_path = os.getenv("NYRA_GGUF_PATH")
    if env_path and Path(env_path).is_file():
        return Path(env_path)

    candidates = [GGUF_FILENAME, *FALLBACK_FILENAMES]
    search_roots: List[Path] = []
    for env_key in ("HF_HOME", "HUGGINGFACE_HUB_CACHE", "TRANSFORMERS_CACHE"):
        val = os.getenv(env_key)
        if val:
            search_roots.append(Path(val))

    home = Path.home()
    search_roots.extend(
        [
            home / ".cache" / "huggingface" / "hub",
            home / ".cache" / "huggingface",
            Path("/data"),
            Path("/data/.huggingface"),
            ROOT / "models",
        ]
    )

    for name in candidates:
        for root in search_roots:
            if not root.exists():
                continue
            direct = root / name
            if direct.is_file():
                return direct
            try:
                for match in root.rglob(name):
                    if match.is_file():
                        return match
            except OSError:
                continue
    return None


def _download_gguf() -> Path:
    if not should_load_model():
        raise RuntimeError(
            "Model download blocked: not running on Hugging Face Spaces. "
            "Local mode is mock-only."
        )

    from huggingface_hub import hf_hub_download

    last_err: Optional[Exception] = None
    for filename in [GGUF_FILENAME, *FALLBACK_FILENAMES]:
        try:
            path = hf_hub_download(
                repo_id=REPO_ID,
                filename=filename,
                resume_download=True,
            )
            return Path(path)
        except Exception as exc:  # noqa: BLE001
            last_err = exc
            continue
    raise RuntimeError(f"Failed to download GGUF from {REPO_ID}: {last_err}")


def _pip_install(args: List[str], env: Optional[dict] = None) -> None:
    cmd = [sys.executable, "-m", "pip", "install", "--quiet", *args]
    subprocess.check_call(cmd, env=env or os.environ.copy())


def _try_import_llama() -> bool:
    try:
        _preload_cuda_from_torch()
        import llama_cpp  # noqa: F401

        return True
    except Exception:
        return False


def _ensure_llama_cpp() -> None:
    """Install llama-cpp-python with CUDA wheel preferred; fix libcudart via torch."""
    global _llama_ready
    if _llama_ready and _try_import_llama():
        return

    _preload_cuda_from_torch()
    if _try_import_llama():
        _llama_ready = True
        return

    env = os.environ.copy()
    torch_lib = _preload_cuda_from_torch()
    if torch_lib:
        env["LD_LIBRARY_PATH"] = os.environ.get("LD_LIBRARY_PATH", "")

    install_attempts = [
        # CUDA 12 wheel (ZeroGPU / modern torch)
        (
            [
                "llama-cpp-python",
                "--force-reinstall",
                "--no-cache-dir",
                "--extra-index-url",
                "https://abetlen.github.io/llama-cpp-python/whl/cu124",
            ],
            env,
        ),
        (
            [
                "llama-cpp-python",
                "--force-reinstall",
                "--no-cache-dir",
                "--extra-index-url",
                "https://abetlen.github.io/llama-cpp-python/whl/cu121",
            ],
            env,
        ),
        # Last resort: default wheel / sdist
        (["llama-cpp-python", "--force-reinstall", "--no-cache-dir"], env),
    ]

    last: Optional[Exception] = None
    for args, e in install_attempts:
        try:
            _pip_install(args, env=e)
            _preload_cuda_from_torch()
            if _try_import_llama():
                _llama_ready = True
                return
        except Exception as exc:  # noqa: BLE001
            last = exc
            continue

    raise RuntimeError(
        "Failed to install/import llama-cpp-python "
        f"(libcudart / CUDA). Last error: {last}"
    )


def get_llm():
    """Singleton Llama. Spaces only (or FORCE_LOCAL_MODEL)."""
    global _llm, _llm_error
    if not should_load_model():
        return None
    if _llm is not None:
        return _llm

    with _llm_lock:
        if _llm is not None:
            return _llm
        try:
            _ensure_llama_cpp()
            _preload_cuda_from_torch()
            from llama_cpp import Llama

            path = _find_gguf_path()
            if path is None:
                path = _download_gguf()

            # Prefer GPU layers when CUDA is visible; fall back to CPU layers=0
            n_gpu = N_GPU_LAYERS
            try:
                import torch

                if not torch.cuda.is_available() and n_gpu != 0:
                    # ZeroGPU may report cuda only inside @spaces.GPU;
                    # still try n_gpu layers — llama.cpp uses its own CUDA.
                    pass
            except ImportError:
                pass

            _llm = Llama(
                model_path=str(path),
                n_ctx=N_CTX,
                n_gpu_layers=n_gpu,
                chat_format="chatml",
                verbose=False,
                logits_all=False,
            )
            _llm_error = None
            return _llm
        except Exception as exc:  # noqa: BLE001
            # Retry once with CPU-only layers if GPU load fails
            try:
                _ensure_llama_cpp()
                from llama_cpp import Llama

                path = _find_gguf_path() or _download_gguf()
                _llm = Llama(
                    model_path=str(path),
                    n_ctx=min(N_CTX, 2048),
                    n_gpu_layers=0,
                    chat_format="chatml",
                    verbose=False,
                )
                _llm_error = f"GPU load failed ({exc}); running CPU fallback"
                return _llm
            except Exception as exc2:  # noqa: BLE001
                _llm_error = str(exc2)
                raise RuntimeError(
                    f"Model load failed. GPU err: {exc} | CPU err: {exc2}"
                ) from exc2


def history_to_messages(
    history: Iterable,
    thinking: bool = False,
) -> List[dict]:
    messages: List[dict] = [
        {"role": "system", "content": load_system_prompt(thinking=thinking)}
    ]

    for item in history or []:
        if isinstance(item, dict):
            role = item.get("role")
            content = item.get("content", "")
            if role in {"user", "assistant"} and content is not None:
                # Skip empty assistant placeholders
                if role == "assistant" and not str(content).strip():
                    continue
                messages.append({"role": role, "content": str(content)})
            continue

        if isinstance(item, (list, tuple)) and len(item) >= 2:
            user, assistant = item[0], item[1]
            if user:
                messages.append({"role": "user", "content": str(user)})
            if assistant:
                messages.append(
                    {"role": "assistant", "content": str(assistant)}
                )

    return messages


def _detect_lang(text: str) -> str:
    t = (text or "").lower()
    pt_signals = [
        "você",
        "voce",
        "olá",
        "ola",
        "obrigado",
        "por que",
        "porque",
        "não",
        "nao",
        "como",
        "está",
        "esta",
        "quero",
        "faça",
        "faca",
        "ajuda",
        "explique",
        "oi",
    ]
    if any(s in t for s in pt_signals) or re.search(
        r"[áàâãéêíóôõúç]", t, re.I
    ):
        return "pt"
    return "en"


def mock_stream(user_text: str) -> Generator[str, None, None]:
    lang = _detect_lang(user_text)
    if lang == "pt":
        reply = (
            "Oi — eu sou a **Nyra**. "
            "Aqui no ambiente local o chat roda em **modo demo** "
            "(sem baixar o modelo). "
            "No **Hugging Face Spaces** com ZeroGPU, a Nyra carrega o "
            "Qwen3.6-35B-A3B (GGUF) e responde de verdade.\n\n"
            f"Você disse: *{user_text[:280]}*"
        )
    else:
        reply = (
            "Hi — I'm **Nyra**. "
            "Locally this UI runs in **demo mode** (no model download). "
            "On **Hugging Face Spaces** with ZeroGPU, Nyra loads "
            "Qwen3.6-35B-A3B (GGUF) and streams real replies.\n\n"
            f"You said: *{user_text[:280]}*"
        )

    acc = ""
    for w in re.split(r"(\s+)", reply):
        acc += w
        yield acc
        time.sleep(0.012)


def stream_chat(
    history: list,
    user_text: str,
    temperature: float = DEFAULT_TEMPERATURE,
    max_tokens: int = DEFAULT_MAX_TOKENS,
    thinking: bool = False,
) -> Generator[str, None, None]:
    """Yield cumulative assistant text."""
    user_text = (user_text or "").strip()
    if not user_text:
        return

    if not should_load_model():
        yield from mock_stream(user_text)
        return

    lang = _detect_lang(user_text)

    try:
        # Ensure GGUF is on disk before Llama() (download is CPU-side)
        prep_err = ensure_weights_ready()
        if prep_err:
            raise RuntimeError(prep_err)
        llm = get_llm()
    except Exception as exc:  # noqa: BLE001
        if lang == "pt":
            yield (
                f"❌ Não consegui carregar o modelo: `{exc}`\n\n"
                "Dica: abra os logs do Space, confirme ZeroGPU e o preload do GGUF. "
                "A primeira carga do Q4 (~21GB) pode falhar se o warmup ainda não terminou — tente de novo em 1–2 min."
            )
        else:
            yield (
                f"❌ Could not load the model: `{exc}`\n\n"
                "Tip: check Space logs, ZeroGPU, and GGUF preload. "
                "First Q4 (~21GB) load may fail if warmup is still running — retry in 1–2 min."
            )
        return

    messages = history_to_messages(history, thinking=thinking)
    messages.append({"role": "user", "content": user_text})

    acc = ""
    kwargs = dict(
        messages=messages,
        temperature=float(temperature),
        max_tokens=int(max_tokens),
        top_p=0.85,
        stream=True,
    )
    try:
        stream = llm.create_chat_completion(
            **kwargs,
            top_k=20,
            presence_penalty=1.2,
            chat_template_kwargs={"enable_thinking": bool(thinking)},
        )
    except TypeError:
        try:
            stream = llm.create_chat_completion(**kwargs, top_k=20)
        except TypeError:
            stream = llm.create_chat_completion(
                messages=messages,
                temperature=float(temperature),
                max_tokens=int(max_tokens),
                stream=True,
            )

    for chunk in stream:
        try:
            delta = chunk["choices"][0]["delta"].get("content") or ""
        except (KeyError, IndexError, TypeError):
            delta = ""
        if delta:
            # Drop boot message once real tokens arrive
            acc += delta
            yield acc

    if not acc:
        yield (
            "_(empty model response)_"
            if lang == "en"
            else "_(resposta vazia do modelo)_"
        )


def status_label(lang: str = "pt") -> str:
    mode = runtime_mode()
    if mode == "spaces":
        return (
            "ZeroGPU · modelo no Space"
            if lang == "pt"
            else "ZeroGPU · model on Space"
        )
    if mode == "local-forced":
        return (
            "Local · modelo forçado"
            if lang == "pt"
            else "Local · forced model"
        )
    return (
        "Demo local · sem download"
        if lang == "pt"
        else "Local demo · no download"
    )


def ensure_weights_ready() -> Optional[str]:
    """
    CPU-side prep (safe outside @spaces.GPU):
    - fix CUDA lib path
    - ensure llama-cpp importable
    - ensure GGUF file is on disk (preload or download)
    Does NOT call Llama() — that stays inside GPU for n_gpu_layers.
    Returns error string or None.
    """
    if not should_load_model():
        return None
    try:
        _preload_cuda_from_torch()
        try:
            import llama_cpp  # noqa: F401
        except Exception:
            _ensure_llama_cpp()
        path = _find_gguf_path()
        if path is None:
            path = _download_gguf()
        return None if path else "GGUF path missing"
    except Exception as exc:  # noqa: BLE001
        return str(exc)


def start_background_warmup() -> None:
    """Kick off CPU-side weight download at Space boot (non-blocking)."""
    global _warmup_started
    if _warmup_started or not should_load_model():
        return
    _warmup_started = True

    def _run():
        err = ensure_weights_ready()
        if err:
            print(f"[nyra] warmup warning: {err}", flush=True)
        else:
            print("[nyra] warmup OK — GGUF + llama-cpp ready", flush=True)

    threading.Thread(target=_run, name="nyra-warmup", daemon=True).start()