"""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 available in llama-cpp-python
    if settings["n_gpu_layers"] != 0:
        try:
            from llama_cpp import Llama
            # Quick test — if CUDA not compiled, will fall back
            test_supported = True
        except Exception:
            test_supported = False

        if not test_supported:
            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

    if len(sys.argv) >= 2:
        if sys.argv[1] == "--text":
            text_only = True
        else:
            image_path = sys.argv[1]

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

    from llama_cpp import Llama

    if text_only:
        # Auto-select: 8GB+ VRAM → Qwen3-8B (smarter), <8GB → Qwen2.5-7B (faster)
        models_base = os.path.expanduser("~/.muravey-china/models")
        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")

        vram_mb, _ = _detect_gpu()
        if vram_mb >= 8000 and os.path.isfile(qwen3_path):
            text_model_path = qwen3_path
            sys.stderr.write(f"[worker] Using Qwen3-8B (VRAM {vram_mb}MB >= 8GB)\n")
        elif os.path.isfile(qwen25_path):
            text_model_path = qwen25_path
            sys.stderr.write(f"[worker] Using Qwen2.5-7B (VRAM {vram_mb}MB < 8GB)\n")
        else:
            text_model_path = qwen3_path  # fallback to whatever exists
        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)
        model = Llama(
            model_path=text_model_path,
            chat_format="chatml",
            n_ctx=n_ctx,
            n_gpu_layers=n_gpu,
            verbose=False,
        )
        # Disable Qwen3 thinking mode — use direct answering (faster, no <think> truncation)
        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)

        handler = Qwen25VLChatHandler(clip_model_path=mmproj_path)
        # Vision model always on CPU — GPU causes encoding bugs in 0.3.33
        model = Llama(
            model_path=model_path,
            chat_handler=handler,
            n_ctx=n_ctx,
            n_gpu_layers=0,
            verbose=False,
        )

        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},
                ],
            }
        ]

    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()
