import os
import re
import gradio as gr
from transformers import DonutProcessor, VisionEncoderDecoderModel
import torch
import traceback

# 1) Load pretrained Donut model and processor
MODEL_NAME = "naver-clova-ix/donut-base-finetuned-cord-v2"
processor = DonutProcessor.from_pretrained(MODEL_NAME)
model = VisionEncoderDecoderModel.from_pretrained(MODEL_NAME)

# 2) Set device and move model
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)

# 3) Inference function with debugging
def ocr_donut(image):
    try:
        if image is None:
            return {"error": "No image provided."}

        # Prepare prompt and inputs
        task_prompt = "<s_cord-v2>"
        decoder_input_ids = processor.tokenizer(
            task_prompt,
            add_special_tokens=False,
            return_tensors="pt"
        ).input_ids.to(device)

        # Convert to tensor
        pixel_values = processor(image.convert("RGB"), return_tensors="pt").pixel_values.to(device)

        # Generate outputs
        outputs = model.generate(
            pixel_values,
            decoder_input_ids=decoder_input_ids,
            max_length=model.config.decoder.max_position_embeddings,
            pad_token_id=processor.tokenizer.pad_token_id,
            eos_token_id=processor.tokenizer.eos_token_id,
            use_cache=True,
            bad_words_ids=[[processor.tokenizer.unk_token_id]],
            return_dict_in_generate=True,
        )

        # Decode and clean up
        sequence = processor.batch_decode(outputs.sequences)[0]
        sequence = sequence.replace(processor.tokenizer.eos_token, "").replace(processor.tokenizer.pad_token, "")
        sequence = re.sub(r"<.*?>", "", sequence, count=1).strip()
        json_output = processor.token2json(sequence)

        return {"result": json_output}

    except Exception:
        tb = traceback.format_exc()
        print(tb)
        return {"error": tb}

# 4) Build Gradio interface
demo = gr.Interface(
    fn=ocr_donut,
    inputs=gr.Image(type="pil", label="Upload Document Image"),
    outputs=gr.JSON(label="Output"),
    title="Donut OCR Gradio App",
    description="Upload a document image and get structured JSON output. Errors will be shown for debugging."
)

# 5) Launch for Spaces
demo.launch(
    server_name="0.0.0.0",
    server_port=int(os.environ.get("PORT", 7860)),
    debug=True
)