"""Qwen2.5-VL worker — runs in separate process to avoid DLL conflicts.

Auto-detects GPU VRAM and adjusts settings:
  8GB+ VRAM: full GPU, n_ctx=2048, max_tokens=256
  4-8GB VRAM: partial GPU, n_ctx=2048, max_tokens=256
  <4GB or no GPU: CPU only, n_ctx=2048, max_tokens=256
"""
import base64
import json
import os
import sys

os.environ["PYTHONIOENCODING"] = "utf-8"
os.environ["HF_HUB_OFFLINE"] = "1"
if hasattr(sys.stdout, "reconfigure"):
    sys.stdout.reconfigure(encoding="utf-8")
if hasattr(sys.stdin, "reconfigure"):
    sys.stdin.reconfigure(encoding="utf-8")


def _detect_gpu():
    """Detect GPU VRAM in MB. Returns (vram_mb, gpu_name) or (0, None)."""
    try:
        import subprocess
        result = subprocess.run(
            ["nvidia-smi", "--query-gpu=memory.total,name", "--format=csv,noheader,nounits"],
            capture_output=True, text=True, timeout=5,
        )
        if result.returncode == 0:
            line = result.stdout.strip().split("\n")[0]
            parts = line.split(",")
            vram = int(parts[0].strip())
            name = parts[1].strip() if len(parts) > 1 else "Unknown"
            return vram, name
    except Exception:
        pass
    return 0, None


def _get_settings():
    """Auto-detect GPU and return optimal settings."""
    vram_mb, gpu_name = _detect_gpu()

    if vram_mb >= 8000:
        # 8GB+ — full GPU
        settings = {
            "n_gpu_layers": -1,
            "n_ctx": 4096,
            "max_tokens": 1024,
            "tier": "HIGH",
        }
    elif vram_mb >= 4000:
        # 4-8GB — partial GPU (RTX 30 series)
        settings = {
            "n_gpu_layers": 20,
            "n_ctx": 2048,
            "max_tokens": 256,
            "tier": "MEDIUM",
        }
    elif vram_mb >= 2000:
        # 2-4GB — only 3B on GPU
        settings = {
            "n_gpu_layers": 10,
            "n_ctx": 2048,
            "max_tokens": 512,
            "tier": "LOW",
        }
    else:
        # No GPU or <2GB — CPU only
        settings = {
            "n_gpu_layers": 0,
            "n_ctx": 2048,
            "max_tokens": 512,
            "tier": "CPU",
        }

    # Check if CUDA is actually available in llama-cpp-python
    if settings["n_gpu_layers"] != 0:
        cuda_ok = False
        try:
            import llama_cpp
            # Check if llama-cpp-python was compiled with CUDA/cuBLAS support
            if hasattr(llama_cpp, "llama_supports_gpu_offload"):
                cuda_ok = llama_cpp.llama_supports_gpu_offload()
            elif hasattr(llama_cpp, "llama_backend_init"):
                # Older versions — try to load with GPU and see
                cuda_ok = True
            else:
                cuda_ok = True  # assume yes if we can't check
        except Exception as exc:
            sys.stderr.write(f"[worker] CUDA check failed: {exc}\n")
            cuda_ok = False

        if not cuda_ok:
            sys.stderr.write("[worker] CUDA not available in llama-cpp-python, falling back to CPU\n")
            settings["n_gpu_layers"] = 0
            settings["tier"] = "CPU"

    sys.stderr.write(f"[worker] GPU: {gpu_name or 'None'} ({vram_mb}MB) → {settings['tier']} mode\n")
    sys.stderr.flush()
    return settings


def main():
    prompt = sys.stdin.read().strip()
    if not prompt:
        prompt = "Describe what you see in this image."

    image_path = None
    text_only = False

    mode = "math"  # default
    if len(sys.argv) >= 2:
        if sys.argv[1] == "--text":
            text_only = True
        elif sys.argv[1] == "--text-chinese":
            text_only = True
            mode = "chinese"
        else:
            image_path = sys.argv[1]
    # Allow explicit --mode flag after image path or --text
    for i, arg in enumerate(sys.argv[2:], 2):
        if arg == "--mode" and i + 1 < len(sys.argv):
            mode = sys.argv[i + 1]

    settings = _get_settings()
    n_gpu = settings["n_gpu_layers"]
    n_ctx = settings["n_ctx"]
    max_tok = settings["max_tokens"]

    try:
        from llama_cpp import Llama
    except Exception as exc:
        sys.stderr.write(f"[worker] FATAL: cannot import llama_cpp: {exc}\n")
        sys.stderr.flush()
        print(json.dumps({"error": f"llama_cpp import failed: {exc}"}))
        sys.exit(1)

    if text_only:
        models_base = os.path.expanduser("~/.muravey-china/models")
        vram_mb, _ = _detect_gpu()

        # Model paths — hybrid selection by mode
        phi4_path   = os.path.join(models_base, "Phi-4-mini-reasoning-Q4_K_M.gguf")
        qwen35_path = os.path.join(models_base, "Qwen3.5-4B-Q4_K_M.gguf")
        qwen3_path  = os.path.join(models_base, "Qwen3-8B-Q4_K_M.gguf")
        qwen25_path = os.path.join(models_base, "qwen2.5-7b-instruct-q3_k_m.gguf")

        is_chinese = mode in ("chinese", "stem", "humanities")

        if is_chinese:
            # Chinese mode: Qwen3.5-4B (best Chinese, 4GB) → Qwen2.5-7B → Qwen3-8B (only if 8GB+)
            if os.path.isfile(qwen35_path):
                text_model_path = qwen35_path
                sys.stderr.write(f"[worker] Chinese mode → Qwen3.5-4B\n")
            elif os.path.isfile(qwen25_path):
                text_model_path = qwen25_path
                sys.stderr.write(f"[worker] Chinese mode → Qwen2.5-7B (fallback)\n")
            elif vram_mb >= 8000 and os.path.isfile(qwen3_path):
                text_model_path = qwen3_path
                sys.stderr.write(f"[worker] Chinese mode → Qwen3-8B (fallback, VRAM {vram_mb}MB)\n")
            else:
                # Last resort: use Qwen3.5-4B path (will error with "not found" below)
                text_model_path = qwen35_path
        else:
            # Math/Physics/Chemistry: Phi-4 → Qwen2.5-7B → Qwen3-8B (only if 8GB+)
            if os.path.isfile(phi4_path):
                text_model_path = phi4_path
                sys.stderr.write(f"[worker] Math mode → Phi-4-mini-Reasoning (MATH 94.6%)\n")
            elif os.path.isfile(qwen25_path):
                text_model_path = qwen25_path
                sys.stderr.write(f"[worker] Math mode → Qwen2.5-7B (fallback)\n")
            elif vram_mb >= 8000 and os.path.isfile(qwen3_path):
                text_model_path = qwen3_path
                sys.stderr.write(f"[worker] Math mode → Qwen3-8B (fallback, VRAM {vram_mb}MB)\n")
            else:
                text_model_path = phi4_path

        if not os.path.isfile(text_model_path):
            sys.stderr.write(f"[worker] ERROR: Model not found: {text_model_path}\n")
            sys.stderr.write("[worker] Please restart program to download models.\n")
            sys.stderr.flush()
            print(json.dumps({"error": "Model not found"}))
            sys.exit(1)

        # Phi-4 uses phi3 chat format; Qwen models use chatml
        is_phi4 = "Phi-4" in os.path.basename(text_model_path)
        is_qwen35 = "Qwen3.5" in os.path.basename(text_model_path)
        chat_fmt = "phi3" if is_phi4 else "chatml"

        model = Llama(
            model_path=text_model_path,
            chat_format=chat_fmt,
            n_ctx=n_ctx,
            n_gpu_layers=n_gpu,
            verbose=False,
        )

        if is_phi4:
            # Phi-4-mini-Reasoning: use thinking mode for better accuracy
            messages = [
                {"role": "system", "content": "You are a helpful assistant specialized in exam questions. Reason step by step, then give a concise final answer."},
                {"role": "user", "content": prompt},
            ]
        elif is_qwen35:
            # Qwen3.5: no_think for speed, direct answer
            messages = [
                {"role": "system", "content": "You are a helpful assistant. Answer directly without using <think> tags."},
                {"role": "user", "content": prompt + "\n\n/no_think"},
            ]
        else:
            # Qwen3-8B / Qwen2.5-7B fallback
            messages = [
                {"role": "system", "content": "You are a helpful assistant. Answer directly without using <think> tags."},
                {"role": "user", "content": prompt + "\n\n/no_think"},
            ]
    else:
        from llama_cpp.llama_chat_format import Qwen25VLChatHandler

        models_dir = os.path.expanduser("~/.muravey-china/models/qwen25vl")
        model_path = os.path.join(models_dir, "Qwen2.5-VL-3B-Instruct-Q4_K_M.gguf")
        mmproj_path = os.path.join(models_dir, "mmproj-Qwen2.5-VL-3B-Instruct-f16.gguf")

        for fpath, fname in [(model_path, "Qwen2.5-VL-3B"), (mmproj_path, "mmproj")]:
            if not os.path.isfile(fpath):
                sys.stderr.write(f"[worker] ERROR: {fname} not found: {fpath}\n")
                sys.stderr.write("[worker] Please restart program to download models.\n")
                sys.stderr.flush()
                print(json.dumps({"error": f"{fname} not found"}))
                sys.exit(1)

        try:
            handler = Qwen25VLChatHandler(clip_model_path=mmproj_path)
        except Exception as exc:
            sys.stderr.write(f"[worker] FATAL: Qwen25VLChatHandler init failed: {exc}\n")
            sys.stderr.flush()
            print(json.dumps({"error": f"Vision handler init failed: {exc}"}))
            sys.exit(1)

        # Vision model always on CPU — GPU causes encoding bugs in 0.3.33
        try:
            model = Llama(
                model_path=model_path,
                chat_handler=handler,
                n_ctx=n_ctx,
                n_gpu_layers=0,
                verbose=False,
            )
        except Exception as exc:
            sys.stderr.write(f"[worker] FATAL: Llama model load failed: {exc}\n")
            sys.stderr.flush()
            print(json.dumps({"error": f"Model load failed: {exc}"}))
            sys.exit(1)

        with open(image_path, "rb") as f:
            img_b64 = base64.b64encode(f.read()).decode("ascii")
        data_uri = f"data:image/jpeg;base64,{img_b64}"
        messages = [
            {
                "role": "user",
                "content": [
                    {"type": "image_url", "image_url": {"url": data_uri}},
                    {"type": "text", "text": prompt},
                ],
            }
        ]

    # For vision: validate image is not blank/white before sending to model
    if not text_only and image_path:
        try:
            import base64 as _b64
            img_bytes = _b64.b64decode(img_b64) if img_b64 else b""
            if len(img_bytes) < 5000:
                # Image too small — likely blank/white screenshot
                print(json.dumps({"error": "BLANK_SCREENSHOT", "text": "NOMATH"}))
                sys.exit(0)
        except Exception:
            pass

    if text_only:
        # Text solver (Qwen3-8B): NO repeat_penalty — math needs repeated variables
        result = model.create_chat_completion(
            messages=messages,
            max_tokens=max_tok,
            temperature=0.1,
        )
    else:
        # Vision OCR (Qwen2.5-VL-3B): repeat_penalty stops reading same text 10+ times
        # 1.1 is gentle — won't hurt formula accuracy but stops looping
        result = model.create_chat_completion(
            messages=messages,
            max_tokens=max_tok,
            temperature=0.1,
            repeat_penalty=1.1,
        )

    text = result["choices"][0]["message"]["content"] or ""
    print(json.dumps({"text": text.strip()}, ensure_ascii=False))


if __name__ == "__main__":
    main()
