import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # noqa: E402  (must come before torch / transformers)
import torch  # noqa: E402
import gradio as gr  # noqa: E402
from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor  # noqa: E402
from qwen_vl_utils import process_vision_info  # noqa: E402

# =============================================================================
# Web-CogReasoner — a Qwen2.5-VL-7B based multimodal web agent that performs
# knowledge-driven Chain-of-Thought reasoning over webpage screenshots.
# https://huggingface.co/Gnonymous/Web-CogReasoner  (paper: arXiv:2508.01858)
# =============================================================================
MODEL_ID = "Gnonymous/Web-CogReasoner"

model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.bfloat16,
    attn_implementation="sdpa",
).to("cuda")
model.eval()

# Cap image tokens so a full-page screenshot doesn't blow up the sequence.
processor = AutoProcessor.from_pretrained(
    MODEL_ID,
    min_pixels=256 * 28 * 28,
    max_pixels=1280 * 28 * 28,
)

# The model was trained as a cognitive web agent: given a webpage screenshot and
# a task, it reasons step-by-step (grounded in factual / conceptual / procedural
# knowledge) before deciding what to do next.
DEFAULT_SYSTEM = (
    "You are Web-CogReasoner, a cognitive web agent. Given a webpage screenshot "
    "and a user task, reason step-by-step about the page and the task before "
    "concluding. First think through the relevant knowledge and the current "
    "state of the page, then state the single next action to take."
)


@spaces.GPU(duration=120)
def analyze(
    image_path: str,
    task: str,
    system_prompt: str = DEFAULT_SYSTEM,
    max_new_tokens: int = 640,
    temperature: float = 0.0,
) -> str:
    """Reason about a webpage screenshot for a given task.

    Args:
        image_path: path to the webpage screenshot to analyze.
        task: the natural-language task / instruction the agent should reason about.
        system_prompt: the system instruction that frames the agent's behaviour.
        max_new_tokens: maximum number of tokens to generate.
        temperature: sampling temperature; 0 uses greedy decoding.

    Returns:
        The model's chain-of-thought reasoning and proposed next action.
    """
    if image_path is None:
        return "Please upload a webpage screenshot first."
    if not task or not task.strip():
        task = "Describe this webpage and suggest the next useful action."

    file_uri = f"file://{os.path.abspath(image_path)}"
    messages = []
    if system_prompt and system_prompt.strip():
        messages.append({"role": "system", "content": system_prompt.strip()})
    messages.append(
        {
            "role": "user",
            "content": [
                {"type": "image", "image": file_uri},
                {"type": "text", "text": task.strip()},
            ],
        }
    )

    chat_text = processor.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    image_inputs, video_inputs = process_vision_info(messages)
    inputs = processor(
        text=[chat_text],
        images=image_inputs,
        videos=video_inputs,
        padding=True,
        return_tensors="pt",
    ).to("cuda")

    do_sample = temperature and temperature > 0
    gen_ids = model.generate(
        **inputs,
        max_new_tokens=int(max_new_tokens),
        do_sample=bool(do_sample),
        temperature=float(temperature) if do_sample else None,
    )
    trimmed = [out[len(inp):] for inp, out in zip(inputs["input_ids"], gen_ids)]
    output = processor.batch_decode(
        trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
    )[0]
    return output.strip()


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

INTRO = """
# 🕸️ Web-CogReasoner

Knowledge-induced **cognitive reasoning** for web agents (Qwen2.5-VL-7B based).
Upload a **webpage screenshot** and describe a **task** — the model reasons
step-by-step about the page and proposes the next action.

[Model](https://huggingface.co/Gnonymous/Web-CogReasoner) ·
[Paper (arXiv:2508.01858)](https://arxiv.org/abs/2508.01858) ·
[Code](https://github.com/Gnonymous/Web-CogReasoner)
"""

with gr.Blocks() as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(INTRO)
        with gr.Row():
            with gr.Column(scale=1):
                image = gr.Image(label="Webpage screenshot", type="filepath")
                task = gr.Textbox(
                    label="Task / instruction",
                    placeholder="e.g. Find and open the login page.",
                    lines=2,
                )
                run = gr.Button("Reason", variant="primary")
            with gr.Column(scale=1):
                output = gr.Textbox(
                    label="Cognitive reasoning & next action",
                    lines=20,
                )

        with gr.Accordion("Advanced settings", open=False):
            system_prompt = gr.Textbox(
                label="System prompt", value=DEFAULT_SYSTEM, lines=4
            )
            max_new_tokens = gr.Slider(
                64, 1024, value=640, step=16, label="Max new tokens"
            )
            temperature = gr.Slider(
                0.0, 1.0, value=0.0, step=0.05, label="Temperature (0 = greedy)"
            )

        gr.Examples(
            examples=[
                ["huggingface.png", "Find how to browse the available models."],
                ["wikipedia.png", "Search for the article about machine learning."],
                ["hackernews.png", "Open the comments for the top story."],
            ],
            inputs=[image, task],
            outputs=output,
            fn=analyze,
            cache_examples=True,
            cache_mode="lazy",
        )

    run.click(
        analyze,
        inputs=[image, task, system_prompt, max_new_tokens, temperature],
        outputs=output,
        api_name="analyze",
    )

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