"""
HuggingFace Space: Blackjack screenshot → JSON state extractor + Q&A.

Designed for ZeroGPU (free H200 bursts). Falls back to CPU if `spaces` isn't
available — expect very slow inference on CPU (a 2B vision model).

Models pulled from public repos:
  base  : unsloth/Qwen3.5-2B
  lora  : davidr99/qwen35-2b-blackjack-reasoning-lora-v4
"""
from __future__ import annotations
import json
import os
import re

import gradio as gr
import torch
from PIL import Image
from peft import PeftModel
from transformers import AutoModelForImageTextToText, AutoProcessor

# ZeroGPU: gives the decorated function GPU access on demand.
# On CPU-only Spaces, this is a no-op shim.
try:
    import spaces  # type: ignore
    GPU_DECORATOR = spaces.GPU(duration=120)
except Exception:
    def _noop(fn):
        return fn
    GPU_DECORATOR = _noop

BASE_MODEL = "unsloth/Qwen3.5-2B"
LORA_REPO  = "davidr99/qwen35-2b-blackjack-reasoning-lora-v4"
DEFAULT_INSTRUCTION = (
    "Extract the blackjack game state from this screenshot as a single JSON object."
)
EXAMPLE_PROMPTS = [
    "Extract the blackjack game state from this screenshot as a single JSON object.",
    "What is the dealer's hand?",
    "What is my hand?",
    "What action should I take?",
    "Should I hit or stand?",
    "What is my current bet?",
    "Who won this hand?",
    "Describe what you see in this image.",
    "Have I busted?",
    "Am I winning right now?",
]

# ---------- model load (module top level so it's downloaded/cached once) ----------
print(f"loading base + LoRA…", flush=True)
processor = AutoProcessor.from_pretrained(BASE_MODEL)
base = AutoModelForImageTextToText.from_pretrained(
    BASE_MODEL,
    torch_dtype=torch.float16,
    device_map="cpu",  # ZeroGPU moves to CUDA inside the decorated call
)
model = PeftModel.from_pretrained(base, LORA_REPO)
model.eval()
print("model ready", flush=True)

# ---------- inference ----------
THINK_RE = re.compile(r"<think>(.*?)</think>\s*", re.DOTALL)

def extract_parts(text: str):
    if "<think>" in text:
        m = THINK_RE.search(text)
        think = m.group(1).strip() if m else ""
        after = THINK_RE.sub("", text).strip() if m else text.strip()
    elif "</think>" in text:
        idx = text.index("</think>")
        think = text[:idx].strip()
        after = text[idx + len("</think>"):].strip()
    else:
        think = text.strip()
        after = ""
    after = after.replace("<|im_end|>", "").replace("<|endoftext|>", "").strip()
    start, end = after.find("{"), after.rfind("}")
    raw_json = after[start:end+1] if start != -1 and end > start else after
    try:
        parsed = json.loads(raw_json) if raw_json else None
    except json.JSONDecodeError:
        parsed = None
    return think, parsed, after


@GPU_DECORATOR
@torch.inference_mode()
def run(image: Image.Image | None, instruction: str, max_new_tokens: int = 768):
    if image is None:
        return "no image", "", "", ""
    if image.mode != "RGB":
        image = image.convert("RGB")
    instruction = (instruction or "").strip() or DEFAULT_INSTRUCTION

    # Move to GPU on demand (no-op on CPU)
    if torch.cuda.is_available():
        model.to("cuda")
        device = "cuda"
    else:
        device = "cpu"

    msgs = [{"role": "user", "content": [
        {"type": "image", "image": image},
        {"type": "text",  "text": instruction},
    ]}]
    prompt = processor.apply_chat_template(msgs, add_generation_prompt=True, enable_thinking=True, tokenize=False)
    inputs = processor(text=prompt, images=[image], return_tensors="pt").to(device)
    out = model.generate(
        **inputs,
        max_new_tokens=max_new_tokens,
        do_sample=False,
        temperature=0.0,
    )
    decoded = processor.tokenizer.decode(
        out[0][inputs["input_ids"].shape[1]:],
        skip_special_tokens=False,
    )
    think, parsed, raw_after = extract_parts(decoded)
    if parsed is not None:
        pretty = json.dumps(parsed, indent=2)
    else:
        pretty = raw_after  # for QA prompts the after-text IS the answer
    return decoded, think, pretty, raw_after

# ---------- UI ----------
with gr.Blocks(title="Blackjack state extractor") as demo:
    gr.Markdown(
        "# Blackjack screenshot → JSON state + Q&A\n"
        "Fine-tuned **Qwen3.5-2B** vision model "
        "([LoRA](https://huggingface.co/davidr99/qwen35-2b-blackjack-reasoning-lora-v4)) "
        "trained on a custom blackjack-web-app dataset. "
        "Default prompt extracts a structured JSON game state. "
        "Try the Q&A prompts below to ask natural-language questions about the screenshot."
    )
    with gr.Row():
        with gr.Column():
            img_in = gr.Image(type="pil", label="Drop a blackjack screenshot")
            prompt_in = gr.Textbox(
                label="Prompt",
                value=DEFAULT_INSTRUCTION,
                lines=4,
            )
            with gr.Accordion("Quick prompts", open=False):
                for p in EXAMPLE_PROMPTS:
                    btn = gr.Button(p, size="sm")
                    btn.click(lambda v=p: v, outputs=prompt_in)
            tokens = gr.Slider(64, 2048, value=768, step=32, label="max new tokens")
            go = gr.Button("Run", variant="primary")
        with gr.Column():
            think_out = gr.Textbox(label="Reasoning (<think>)", lines=10)
            answer_out = gr.Code(label="Answer (JSON if extract task, else text)", language="json")
            raw_out = gr.Textbox(label="Raw output (debug)", lines=6)
    go.click(run, inputs=[img_in, prompt_in, tokens],
             outputs=[raw_out, think_out, answer_out, raw_out])

if __name__ == "__main__":
    demo.queue().launch()