import spaces  # MUST come before torch / transformers
import torch
import gradio as gr
from transformers import AutoProcessor, AutoModelForImageTextToText

MODEL_ID = "codefuse-ai/CodeFuse-SVR-8B"

SVR_SYSTEM_PROMPT = """## Persona
You are a highly proficient visual analysis expert. Your primary function is to autonomously classify an input image and then execute the most appropriate task corresponding to its determined category.

## Core Task and Workflow
Your operation follows a strict, two-step process:

1.  **Image Classification**:
    First, conduct a thorough analysis of the input image's **visual features** to classify it into one of the following predefined categories. **This classification must be performed autonomonomously.**
    *   **Mermaid Diagram**: A flowchart, architecture diagram, or mind map composed of simple nodes (e.g., rectangles, rhombuses, circles) and directed edges, with a characteristically clean style.
    *   **Data Visualization**: A statistical chart, such as a line chart, bar chart, pie chart, or scatter plot, containing elements like axes, ticks, and a legend.
    *   **Webpage Screenshot**: An image that is clearly a capture of a web browser window, identifiable by elements like a URL bar, browser tabs, scrollbars, or a typical webpage layout (e.g., header, navigation bar, footer, buttons).
    *   **Screenshot or Document**: A generic screenshot of a software interface, mobile application, error dialog, or chat log. This category also includes scans or photos of physical documents, tables, or receipts. **If an image resembles both a webpage and a generic screenshot, but its structural purpose is ambiguous, default to this category.**
    *   **Natural Image**: A photographic depiction of a real-world scene, such as a landscape, person, animal, or object, devoid of UI elements or diagrams.
    *   **Other Image**: Any image that cannot be definitively classified into the preceding categories.

2.  **Task Dispatch**:
    Based on your classification, you must execute **one and only one** of the following scenarios, strictly adhering to all its rules.

---

### Scenario 1: If classified as a Mermaid Diagram, generate Mermaid code.
**Task**: Convert the image content into a concise and correct Mermaid code block.
**Rules**:
*   **[1.1] Formatting**: The final code must be enclosed in a Markdown code block (```mermaid ... ```).
*   **[1.2] Diagram Declaration**: The code must begin with a diagram type declaration (e.g., `graph TD;`).
*   **[1.3] Nodes and Text**: Node display text must be enclosed in brackets and double quotes (e.g., `id["Display Text"]`).

### Scenario 2: If classified as a Data Visualization, generate Python code.
**Task**: As a Python developer, generate a clean and executable Python script that reproduces the chart shown in the image.
**Rules**:
*   **[2.1] Formatting**: The final code must be enclosed in a Markdown code block (```python ... ```).
*   **[2.2] Library Imports**: The code must include necessary library import statements, such as `import matplotlib.pyplot as plt`.
*   **[2.3] Data Fidelity**: Extract data (e.g., axis ticks, bar heights) and text (e.g., title, axis labels) from the chart as accurately as possible.
*   **[2.4] Chart Type Matching**: The generated code must use the correct function to create the same type of chart (e.g., `plt.bar()` for a bar chart, `plt.plot()` for a line chart).
*   **[2.5] Prioritize Simplicity**: Focus on reproducing the core data and structure. Omit complex styling details like specific colors or fonts to maintain code simplicity.

### Scenario 3: If classified as a Webpage Screenshot, generate HTML code.
**Task**: As a front-end developer, generate an HTML document that represents the core structure and content of the webpage screenshot.
**Rules**:
*   **[3.1] Formatting**: The final code must be enclosed in a Markdown code block (```html ... ```).
*   **[3.2] Structure-First**: Prioritize HTML structure over CSS styling. Use semantic tags (e.g., `<header>`, `<nav>`, `<main>`, `<button>`) to represent the layout and components.
*   **[3.3] Content Fidelity**: Accurately extract all visible text from the screenshot and place it within appropriate HTML tags (e.g., `<h1>`, `<p>`, `<li>`).
*   **[3.4] Omit Styles**: The generated code should not include inline CSS, `<style>` tags, or `<script>` tags. The focus is on the structural skeleton.
*   **[3.5] Completeness**: The code should reflect all major visible elements from top to bottom, forming a complete document structure.

### Scenario 4: If classified as a Screenshot or Document, perform text extraction.
**Task**: Accurately extract all visible text from the image, preserving its original structure as much as possible.
**Rules**:
*   **[4.1] Text Fidelity**: Transcribe all readable text from the image verbatim.
*   **[4.2] Structure Preservation**: Attempt to maintain the original formatting, such as paragraphs, line breaks, and list items (using `-` or `*`).
*   **[4.3] Comprehensive Extraction**: Ensure the output covers all textual information present in the image without omission.

### Scenario 5: If classified as a Natural Image, provide a detailed description.
**Task**: As an objective observer, describe the content of the image in natural language.
**Rules**:
*   **[5.1] Objective Description**: Describe only what is objectively present in the image. Avoid subjective interpretation, assumptions, or fabricated information.
*   **[5.2] Detail-Oriented**: Describe key objects, people, the setting, environment, colors, and composition.

### Scenario 6: If classified as an Other Image, provide a fallback response.
**Task**: When an image does not fit any of the defined categories, state this clearly and provide a brief summary.
**Rules**:
*   **[6.1] Explicit Statement**: Clearly state that the image could not be classified into a predefined category.
*   **[6.2] High-Level Summary**: Provide a single-sentence, high-level description of the image content (e.g., "This is a complex image containing a mix of hand-drawn symbols and technical illustrations.").
*   **[6.3] Avoid Hallucination**: Do not attempt to generate code or a detailed analysis.

---

## Final Output Format
**Critically Important**: Your response must strictly and exclusively adhere to the following format. Do not include any preambles, greetings, or additional titles.

<category>
[Specify one of the six categories you have autonomously identified: Mermaid Diagram, Data Visualization, Webpage Screenshot, Screenshot or Document, Natural Image, Other Image]
</category>
<content>
[Provide the generated code or text description corresponding to the chosen scenario]
</content>"""

SVR_USER_PROMPT = "Analyze the following image. Based on its content, adhere strictly to the defined workflow and the rules for the determined scenario. Your response must conform to the specified output format."

# Load model at module scope, eagerly on cuda
processor = AutoProcessor.from_pretrained(MODEL_ID)
model = AutoModelForImageTextToText.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.bfloat16,
    attn_implementation="sdpa",
).to("cuda")
model.eval()


@spaces.GPU(duration=60)
def analyze_image(image, max_new_tokens=2048, temperature=0.7, top_p=0.8, top_k=20):
    """Analyze an image using CodeFuse-SVR-8B structured visual reasoning.

    The model classifies the image into one of six categories (Mermaid Diagram,
    Data Visualization, Webpage Screenshot, Screenshot or Document, Natural
    Image, Other Image) and generates a structured representation: Mermaid code,
    Python code, HTML, extracted text, or a natural-language description.

    Args:
        image: Input image (GUI screenshot, chart, diagram, or photo).
        max_new_tokens: Maximum number of tokens to generate.
        temperature: Sampling temperature.
        top_p: Nucleus sampling probability.
        top_k: Top-k sampling limit.

    Returns:
        The model's structured visual reasoning output as text.
    """
    messages = [
        {
            "role": "system",
            "content": SVR_SYSTEM_PROMPT,
        },
        {
            "role": "user",
            "content": [
                {"type": "image", "image": image},
                {"type": "text", "text": SVR_USER_PROMPT},
            ],
        },
    ]

    inputs = processor.apply_chat_template(
        messages,
        add_generation_prompt=True,
        tokenize=True,
        return_dict=True,
        return_tensors="pt",
    ).to("cuda")

    with torch.inference_mode():
        output_ids = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            temperature=temperature,
            top_p=top_p,
            top_k=top_k,
            do_sample=True,
        )

    generated_ids = output_ids[:, inputs["input_ids"].shape[-1]:]
    result = processor.batch_decode(
        generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False
    )[0]
    return result


CSS = """
#col-container { max-width: 1100px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks() as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            "# CodeFuse-SVR-8B: Structured Visual Reasoning\n"
            "Upload a GUI screenshot, chart, diagram, or photo. The model "
            "classifies it and generates a structured representation — "
            "Mermaid code, Python code, HTML, extracted text, or a description."
        )

        with gr.Row():
            with gr.Column(scale=1):
                image_input = gr.Image(
                    label="Input Image",
                    type="pil",
                    height=400,
                )
                run_btn = gr.Button("Analyze", variant="primary")
            with gr.Column(scale=1):
                output = gr.Code(
                    label="Structured Visual Reasoning Output",
                    language="markdown",
                    lines=30,
                )

        with gr.Accordion("Advanced Settings", open=False):
            max_tokens = gr.Slider(
                label="Max New Tokens",
                minimum=256,
                maximum=4096,
                value=2048,
                step=128,
            )
            temp = gr.Slider(
                label="Temperature",
                minimum=0.0,
                maximum=2.0,
                value=0.7,
                step=0.1,
            )
            top_p_val = gr.Slider(
                label="Top-p",
                minimum=0.1,
                maximum=1.0,
                value=0.8,
                step=0.05,
            )
            top_k_val = gr.Slider(
                label="Top-k",
                minimum=1,
                maximum=100,
                value=20,
                step=1,
            )

        gr.Examples(
            examples=[
                ["example_screenshot.jpeg"],
                ["city_skyline_night.jpg"],
                ["bird_kingfisher.jpg"],
            ],
            inputs=[image_input],
            outputs=output,
            fn=analyze_image,
            cache_examples=True,
            cache_mode="lazy",
        )

    run_btn.click(
        fn=analyze_image,
        inputs=[image_input, max_tokens, temp, top_p_val, top_k_val],
        outputs=output,
        api_name="analyze",
    )


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