# ruff: noqa: I001

import spaces

import json

import torch
from transformers import AutoModelForImageTextToText, AutoProcessor


processor = AutoProcessor.from_pretrained(
    "numind/NuExtract3",
    trust_remote_code=True,
)

model = (
    AutoModelForImageTextToText.from_pretrained(
        "numind/NuExtract3",
        attn_implementation="sdpa",
        dtype=torch.bfloat16,
        trust_remote_code=True,
    )
    .to("cuda")
    .eval()
)


@spaces.GPU
def extract(image, text, template, enable_thinking):
    inputs = processor.apply_chat_template(
        [
            {
                "role": "user",
                "content": [
                    {"type": "image", "image": image},
                    {"type": "text", "text": text},
                ],
            }
        ],
        add_generation_prompt=True,
        tokenize=True,
        return_dict=True,
        return_tensors="pt",
        mode="structured",
        template=template,
        enable_thinking=enable_thinking,
    ).to(model.device)

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

    output = processor.batch_decode(
        generated_ids[:, inputs.input_ids.shape[1] :],
        skip_special_tokens=True,
        clean_up_tokenization_spaces=False,
    )[0].strip()

    return json.loads(
        output.split("</think>", 1)[1].strip() if enable_thinking else output
    )


@spaces.GPU
def generate_template(image):
    inputs = processor.apply_chat_template(
        [
            {
                "role": "user",
                "content": [
                    {"type": "image", "image": image},
                    {
                        "type": "text",
                        "text": (
                            "Create a reusable structured extraction template grounded only "
                            "in the visible document. Include fields supported by the document, "
                            "represent repeated records as arrays, use NuExtract template leaf "
                            "types, and return only the JSON template."
                        ),
                    },
                ],
            }
        ],
        add_generation_prompt=True,
        tokenize=True,
        return_dict=True,
        return_tensors="pt",
        mode="template-generation",
    ).to(model.device)

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

    return json.dumps(
        json.loads(
            processor.batch_decode(
                generated_ids[:, inputs.input_ids.shape[1] :],
                skip_special_tokens=True,
                clean_up_tokenization_spaces=False,
            )[0].strip()
        ),
        indent=2,
    )