"""Qwen2.5-VL Smart Document, Blueprint & Invoice OCR Studio on ZeroGPU."""

import os

os.environ.setdefault("PYTORCH_ALLOC_CONF", "backend:cudaMallocAsync")
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "backend:cudaMallocAsync")

import spaces  # Must be imported before torch
import gradio as gr
import PIL.Image
import torch
from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration

MODEL_ID = "Qwen/Qwen2.5-VL-3B-Instruct"

EXTRACTION_PRESETS = {
    "Full Clean Markdown + Tables (Default)": (
        "Extract all text, headings, and tables from this image into clean, well-formatted GitHub Flavored Markdown. "
        "Preserve exact numbers, part numbers, prices, and table columns."
    ),
    "Structured Invoice / Receipt JSON": (
        "Extract all fields from this invoice or receipt as valid JSON with keys: vendor, date, invoice_number, "
        "line_items (description, quantity, unit_price, total), subtotal, tax, and total_amount."
    ),
    "Electrical Panel / Blueprint / Spec Analysis": (
        "Analyze this electrical diagram, panel schedule, or technical blueprint. List all visible circuits, "
        "breaker ratings, wire gauges, load calculations, labels, and technical notes in structured Markdown."
    ),
    "Custom Question / Prompt": "",
}

print(f"[vision-ocr-studio] Preloading {MODEL_ID} at module scope for ZeroGPU...", flush=True)
model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.bfloat16,
)
model.to("cuda")
model.eval()
processor = AutoProcessor.from_pretrained(MODEL_ID)
print("[vision-ocr-studio] Qwen2.5-VL OCR engine ready on ZeroGPU!", flush=True)


@spaces.GPU(duration=60)
def analyze_document(
    image_path: str | None,
    preset: str = "Full Clean Markdown + Tables (Default)",
    custom_instruction: str = "",
    max_tokens: int = 1024,
) -> str:
    """Run vision-language OCR or structured extraction on an uploaded document/image."""
    if not image_path:
        raise gr.Error("Please upload a document, blueprint, invoice, or screenshot image first.")

    instruction = (custom_instruction or "").strip() or EXTRACTION_PRESETS.get(
        preset, EXTRACTION_PRESETS["Full Clean Markdown + Tables (Default)"]
    )
    img = PIL.Image.open(image_path).convert("RGB")
    if max(img.size) > 1600:
        scale = 1600.0 / max(img.size)
        img = img.resize((int(img.width * scale), int(img.height * scale)), PIL.Image.Resampling.LANCZOS)

    messages = [
        {
            "role": "user",
            "content": [
                {"type": "image", "image": img},
                {"type": "text", "text": instruction},
            ],
        }
    ]
    text_prompt = processor.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    inputs = processor(
        text=[text_prompt],
        images=[img],
        padding=True,
        return_tensors="pt",
    ).to("cuda")

    with torch.no_grad():
        generated_ids = model.generate(**inputs, max_new_tokens=int(max_tokens))

    trimmed_ids = [
        out_ids[len(in_ids):]
        for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
    ]
    output_text = processor.batch_decode(
        trimmed_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False
    )[0]
    return output_text.strip()


with gr.Blocks(title="Smart Document, Blueprint & Invoice OCR Studio", fill_width=True) as demo:
    gr.Markdown(
        "# 📄 Smart Document, Blueprint & Invoice OCR Studio\n"
        "Powered by **[Qwen/Qwen2.5-VL-3B-Instruct](https://huggingface.co/Qwen/Qwen2.5-VL-3B-Instruct)** on **ZeroGPU (`zero-a10g`)** at **$0/hour**—extract clean **Markdown, JSON, tables, invoices, or electrical specs** from any image."
    )
    with gr.Row():
        with gr.Column(scale=1):
            img_in = gr.Image(
                label="Upload Document, Invoice, Electrical Schedule, Blueprint, or Screenshot",
                type="filepath",
            )
            preset_dd = gr.Dropdown(
                choices=list(EXTRACTION_PRESETS),
                value="Full Clean Markdown + Tables (Default)",
                label="Extraction Mode",
            )
            custom_box = gr.Textbox(
                label="Custom Instruction / Question (Optional — overrides preset when filled)",
                lines=3,
                placeholder="e.g., Extract all circuit breaker amps and wire sizes into a JSON array...",
            )
            max_tok = gr.Slider(256, 2048, value=1024, step=128, label="Max Output Tokens")
            extract_btn = gr.Button("📄 Extract & Analyze Document (ZeroGPU)", variant="primary")
        with gr.Column(scale=1):
            md_out = gr.Textbox(
                label="Extracted Output (Markdown / JSON)",
                lines=18,
                buttons=["copy"],
            )

    extract_btn.click(
        fn=analyze_document,
        inputs=[img_in, preset_dd, custom_box, max_tok],
        outputs=md_out,
        api_name="extract_document",
    )

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Ocean(), mcp_server=True)