"""hitech-vlm — Qwen3-VL-8B-Instruct multimodal Space (vision + text).

Serves the contract that `core/clients.py::_predict_default` expects:

    predict(prompt: str, schema_json: str, image_path: str | None) -> str   # JSON

The composed `prompt` already carries the task instruction *and* the target JSON
schema (see `core/clients.py::_compose_prompt`), so `schema_json` is passed only
to satisfy the 3-arg contract — this Space does not use it separately.

Model choice: the original Qwen3.6-35B-A3B-FP8 deployed and loaded but its
fine-grained-FP8 inference is broken on the current ZeroGPU image — transformers
routes the FP8 matmul to the `kernels-community/deep-gemm` kernel, which rejects
the checkpoint's layout with `RuntimeError: Unknown recipe`. Qwen3.5-FP8 shares
that path, so the pivot is to a quant that needs no special kernels:
**Qwen3-VL-8B-Instruct in bf16** (~16 GB, fits ZeroGPU `large` 48 GB with room to
spare; mainline `qwen3_vl` arch, no `trust_remote_code`, no FP8/AWQ kernels).
Capability upgrade path if Step-6 evals show 8B is too weak on drawings:
`Qwen3-VL-30B-A3B-Instruct-AWQ` (MoE, ~18 GB — but reintroduces AWQ-kernel risk).

ZeroGPU notes (see huggingface-zerogpu skill): eager module-scope load; `import
spaces` is unconditional and omitted from requirements.txt; greedy decode for
clean JSON.
"""

from __future__ import annotations

import json
import re

import cv2
import gradio as gr
import numpy as np
import spaces
import torch
from PIL import Image
from transformers import AutoModelForImageTextToText, AutoProcessor

MODEL_ID = "Qwen/Qwen3-VL-8B-Instruct"
MAX_NEW_TOKENS = 3072

# Eager module-scope load (see ZeroGPU note above). dtype="auto" -> bf16 per the
# checkpoint; mainline qwen3_vl needs no trust_remote_code.
processor = AutoProcessor.from_pretrained(MODEL_ID)
model = AutoModelForImageTextToText.from_pretrained(
    MODEL_ID,
    dtype="auto",
    device_map="auto",
).eval()


def _preprocess_image(pil_img: Image.Image) -> Image.Image:
    """Deskew + auto-rotate a scanned document or photo before VLM inference.

    Steps:
    1. Convert to grayscale and threshold to isolate text/content.
    2. Find the dominant skew angle via Hough lines and rotate to correct it.
    3. Return as RGB PIL Image (unchanged if preprocessing fails for any reason).

    This improves extraction accuracy on scanned POs and hand-marked drawings
    that arrive slightly tilted from a phone camera or flatbed scanner.
    """
    try:
        img = np.array(pil_img.convert("RGB"))
        gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
        _, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)

        # Detect lines to estimate skew angle
        lines = cv2.HoughLinesP(binary, 1, np.pi / 180, threshold=100, minLineLength=100, maxLineGap=10)
        if lines is None or len(lines) == 0:
            return pil_img

        angles = []
        for line in lines:
            x1, y1, x2, y2 = line[0]
            if x2 != x1:
                angles.append(np.degrees(np.arctan2(y2 - y1, x2 - x1)))

        if not angles:
            return pil_img

        # Median angle — robust against outlier lines
        skew = float(np.median(angles))
        # Only correct small skews (> 0.5° and < 45°) to avoid false rotations
        if abs(skew) < 0.5 or abs(skew) > 45:
            return pil_img

        h, w = img.shape[:2]
        center = (w / 2, h / 2)
        M = cv2.getRotationMatrix2D(center, skew, 1.0)
        rotated = cv2.warpAffine(img, M, (w, h), flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_REPLICATE)
        return Image.fromarray(rotated)
    except Exception:
        return pil_img


def _clean_json(text: str) -> str:
    """Best-effort: drop <think> blocks, code fences and prose; keep the JSON object.

    `core/clients.py` does an unforgiving `json.loads` on our return value, so a
    leading fence or preamble would waste the caller's single retry. Each Space
    owns producing a clean JSON string.
    """
    text = re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL)
    fence = re.search(r"```(?:json)?\s*(.*?)```", text, flags=re.DOTALL)
    if fence:
        text = fence.group(1)
    # Return the FIRST complete JSON object. A plain find('{')..rfind('}') span
    # breaks when the model emits two objects (e.g. an example then the answer):
    # the slice would span the gap between them and fail the caller's json.loads.
    start = text.find("{")
    if start != -1:
        try:
            obj, _ = json.JSONDecoder().raw_decode(text[start:])
            return json.dumps(obj)
        except json.JSONDecodeError:
            end = text.rfind("}")  # fallback: original brace-span heuristic
            if end > start:
                return text[start : end + 1].strip()
    return text.strip()


def _build_messages(prompt: str, image_path: str | None) -> list[dict]:
    content: list[dict] = []
    if image_path:
        img = _preprocess_image(Image.open(image_path).convert("RGB"))
        content.append({"type": "image", "image": img})
    content.append({"type": "text", "text": prompt})
    return [{"role": "user", "content": content}]


@spaces.GPU(duration=90)
def predict(prompt: str, schema_json: str, image_path: str | None) -> str:
    """Run one multimodal (or text-only) inference and return a JSON string."""
    messages = _build_messages(prompt, image_path)
    inputs = processor.apply_chat_template(
        messages,
        add_generation_prompt=True,
        tokenize=True,
        return_dict=True,
        return_tensors="pt",
    ).to(model.device)

    with torch.inference_mode():
        generated = model.generate(
            **inputs,
            max_new_tokens=MAX_NEW_TOKENS,
            do_sample=False,
        )

    generated = generated[:, inputs["input_ids"].shape[1] :]
    text = processor.batch_decode(
        generated,
        skip_special_tokens=True,
        clean_up_tokenization_spaces=False,
    )[0]
    return _clean_json(text)


# gr.Interface exposes `fn` at api_name="/predict", which is what core.clients calls.
demo = gr.Interface(
    fn=predict,
    inputs=[
        gr.Textbox(label="prompt"),
        gr.Textbox(label="schema_json"),
        gr.Image(type="filepath", label="image_path"),
    ],
    outputs=gr.Textbox(label="json"),
    title="Hi-Tech VLM",
    description=(
        "Qwen3-VL-8B-Instruct — multimodal JSON extraction for the Hi-Tech AI "
        "Platform. Returns a JSON string for core.clients to Pydantic-validate."
    ),
)

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