import os
os.environ["GRADIO_TEMP_DIR"] = "./tmp"

import time
import torch
import spaces
import tempfile
import sys
import gradio as gr
from io import StringIO
from contextlib import contextmanager
from threading import Thread
from PIL import Image
from transformers import (
    AutoProcessor,
    AutoModelForCausalLM,
    AutoModel,
    AutoTokenizer,
    Qwen2_5_VLForConditionalGeneration,
    TextIteratorStreamer
)
from huggingface_hub import snapshot_download
from qwen_vl_utils import process_vision_info
from otsl_utils import convert_otsl_to_html

# == download weights ==
# model_dir = snapshot_download('opendatalab/TRivia-3B', local_dir='./models/TRivia-3B')
# == select device ==
device = 'cuda' if torch.cuda.is_available() else 'cpu'

# Load TRivia-3B
try:
    MODEL_ID = "opendatalab/TRivia-3B"
    processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
    model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
        MODEL_ID,
        attn_implementation="flash_attention_2",
        trust_remote_code=True,
        torch_dtype=torch.float16,
        device_map="auto"
    ).eval()
    print("✓ TRivia-3B loaded")
except Exception as e:
    model = None
    processor = None

@spaces.GPU
def recognize_image(image: Image.Image,
                   max_new_tokens: int, temperature: float):
    if image is None:
        yield "Please upload an image.", "Please upload an image."
        return

    try:
        # Prepare messages in chat format
        messages = [{
            "role": "user",
            "content": [
                {"type": "text", "text": "You are an AI specialized in recognizing and extracting table from images. Your mission is to analyze the table image and generate the result in OTSL format using specified tags. Output only the results without any other words and explanation."},
                {"type": "image"},
            ]
        }]

        prompt_full = processor.apply_chat_template(
            messages, 
            tokenize=False, 
            add_generation_prompt=True
        )

        inputs = processor(
            text=[prompt_full],
            images=[image],
            return_tensors="pt",
            padding=True
        ).to(device)

        streamer = TextIteratorStreamer(
            processor.tokenizer if hasattr(processor, 'tokenizer') else processor, 
            skip_prompt=True, 
            skip_special_tokens=True
        )


        generation_kwargs = {
            **inputs,
            "streamer": streamer,
            "max_new_tokens": max_new_tokens,
            "do_sample": True,
            "temperature": temperature,
            "repetition_penalty": 1.05,
        }


        thread = Thread(target=model.generate, kwargs=generation_kwargs)
        thread.start()


        # Stream the results
        buffer = ""
        for new_text in streamer:
            buffer += new_text
            buffer = buffer.replace("<|im_end|>", "")
            html_text = convert_otsl_to_html(buffer)
            time.sleep(0.01)
            yield buffer, html_text, html_text


        # Ensure thread completes
        thread.join()


    except Exception as e:
        error_msg = f"Error during generation: {str(e)}"
        print(f"Full error: {e}")
        import traceback
        traceback.print_exc()
        yield error_msg, error_msg, error_msg
    
def gradio_reset():
    return gr.update(value=None), gr.update(value=None), gr.update(value=None), gr.update(value=None)

    
if __name__ == "__main__":
    with open("header.html", "r") as file:
        header = file.read()
    with gr.Blocks() as demo:
        gr.HTML(header)
        
        with gr.Row():
            with gr.Column():
                
                input_img = gr.Image(label=" ", interactive=True)
                with gr.Row():
                    clear = gr.Button(value="Clear")
                    predict = gr.Button(value="Table Recognition", interactive=True, variant="primary")
                    
                with gr.Accordion("Advanced Settings", open=False):
                    max_tokens = gr.Slider(
                        minimum=1,
                        maximum=8192,
                        value=4096,
                        step=1,
                        label="Max New Tokens"
                    )
                    temperature = gr.Slider(
                        minimum=0.1,
                        maximum=2.0,
                        value=0.1,
                        step=0.1,
                        label="Temperature"
                    )
                    
                with gr.Accordion("Examples:"):
                    example_root = os.path.join(os.path.dirname(__file__), "assets", "example")
                    gr.Examples(
                        examples=[os.path.join(example_root, _) for _ in os.listdir(example_root) if
                                    _.endswith("png")],
                        inputs=[input_img],
                    )
            with gr.Column():
                rendered_html = gr.Markdown(label="Rendered HTML:", show_label=True)
                output_html = gr.Textbox(label="Converted HTML:", interactive=False)
                pred_otsl = gr.Textbox(label="Predicted OTSL:", interactive=False)
    
        clear.click(gradio_reset, inputs=None, outputs=[input_img, pred_otsl, output_html, rendered_html])
        predict.click(recognize_image, inputs=[input_img, max_tokens, temperature], outputs=[pred_otsl, output_html, rendered_html])
    
    demo.launch(ssr_mode=False, show_error=True)