import json
import re
import traceback
from pathlib import Path

import gradio as gr
import torch
from PIL import Image
from peft import PeftModel
from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration

# =========================
# CONFIG
# =========================
BASE_MODEL_ID = "Qwen/Qwen2.5-VL-3B-Instruct"

ADAPTER_PATHS = {
    "web": "./adapters/web",
    "mobile": "./adapters/mobile",
}

SYSTEM_PROMPT = """You are an expert UI/UX design analyst. When given a UI screenshot, analyze it thoroughly and return a structured JSON report with exactly these keys:
{
  "ui_elements": [
    {"type": "element type", "label": "text/label if any", "position_hint": "top-left / center / etc."}
  ],
  "layout_structure": "description of overall layout pattern (e.g. sidebar + main, top-nav + grid)",
  "hierarchy": {
    "primary":   ["most prominent / CTA elements"],
    "secondary": ["supporting elements"],
    "tertiary":  ["decorative or minor elements"]
  },
  "style": {
    "color_palette":  ["dominant colors as hex or descriptive names"],
    "typography":     "font style observations",
    "spacing":        "tight / balanced / airy",
    "visual_theme":   "overall aesthetic feel"
  },
  "summary": "one paragraph plain-English summary of the UI"
}
Respond ONLY with valid JSON. No markdown fences, no extra commentary."""


# =========================
# LOAD MODEL
# =========================
print("Loading model...")

try:
    processor = AutoProcessor.from_pretrained(
        BASE_MODEL_ID,
        trust_remote_code=True,
    )

    base_model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
        BASE_MODEL_ID,
        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
        device_map="auto",
        trust_remote_code=True,
    )

    base_model.eval()
    print("Model loaded successfully")

except Exception as e:
    print(f"FATAL: {e}")
    raise

loaded_adapters = {}

# =========================
# ADAPTER LOADER
# =========================
def get_model_with_adapter(adapter_type: str):
    if adapter_type not in ADAPTER_PATHS:
        return base_model

    path = ADAPTER_PATHS[adapter_type]

    if not Path(path).exists():
        return base_model

    if adapter_type in loaded_adapters:
        return loaded_adapters[adapter_type]

    try:
        model = PeftModel.from_pretrained(base_model, path)
        model.eval()
        loaded_adapters[adapter_type] = model
        return model
    except:
        return base_model

# =========================
# HELPERS
# =========================
def clean_json_output(raw: str):
    raw = raw.strip()
    raw = re.sub(r"^```.*?\n", "", raw, flags=re.DOTALL)
    raw = re.sub(r"```$", "", raw)
    return raw.strip()

def parse_json_safe(text: str):
    try:
        return json.loads(text), None
    except Exception as e:
        match = re.search(r"\{.*\}", text, re.DOTALL)
        if match:
            try:
                return json.loads(match.group()), None
            except:
                pass
        return None, str(e)

def build_markdown(parsed, raw, error):
    if parsed is None:
        return f"## ❌ Failed to parse\n\n```\n{raw}\n```"

    return f"""
## ✅ Result

**Screen:** `{parsed.get("screen_id","N/A")}`  
**Summary:** {parsed.get("summary","N/A")}  
**Layout:** `{parsed.get("layout_structure","N/A")}`
"""

# =========================
# MAIN FUNCTION
# =========================
def analyze_ui(image, prompt, adapter_type):

    if image is None:
        return "No image provided", "{}"

    try:
        model = get_model_with_adapter(adapter_type)

        user_prompt = f"Analyze this UI. {prompt or ''}"

        messages = [
            {"role": "system", "content": SYSTEM_PROMPT},
            {
                "role": "user",
                "content": [
                    {"type": "image", "image": image},
                    {"type": "text", "text": user_prompt},
                ],
            },
        ]

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

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

        device = next(model.parameters()).device
        inputs = {k: v.to(device) if hasattr(v, "to") else v for k, v in inputs.items()}

        with torch.no_grad():
            out = model.generate(
                **inputs,
                max_new_tokens=800,
                do_sample=False,
            )

        gen = out[:, inputs["input_ids"].shape[1]:]

        raw = processor.batch_decode(gen, skip_special_tokens=True)[0]

        cleaned = clean_json_output(raw)
        parsed, err = parse_json_safe(cleaned)

        md = build_markdown(parsed, cleaned, err)

        json_out = json.dumps(parsed, indent=2, ensure_ascii=False) if parsed else cleaned

        return md, json_out

    except Exception as e:
        tb = traceback.format_exc()
        return f"## ERROR\n```\n{tb}\n```", "{}"

# =========================
# UI
# =========================
with gr.Blocks() as demo:

    gr.Markdown("# 🖥️ UI Analyzer")

    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="pil")
            adapter = gr.Radio(["web", "mobile"], value="web")
            prompt = gr.Textbox(lines=3)
            btn = gr.Button("Analyze")

        with gr.Column():
            md_out = gr.Markdown()
            json_out = gr.Textbox(lines=25)  # ❌ شيلنا show_copy_button

    btn.click(
        analyze_ui,
        inputs=[image_input, prompt, adapter],
        outputs=[md_out, json_out],
    )

# =========================
# LAUNCH
# =========================
if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)