import gradio as gr
import base64
import io
import json
import os
import re
import urllib.error
import urllib.request
from PIL import Image
# from transformers import AutoModel, AutoProcessor 

# ==========================================
# 1. Configure SiliconFlow API
# ==========================================
API_URL = "https://api.siliconflow.cn/v1/chat/completions"
MODEL_ID = "Qwen/Qwen3.6-35B-A3B"
API_KEY = os.getenv(
    "SILICONFLOW_API_KEY",
    "sk-aflkmenvmawlrzbmlawwnhdyaacpylqeebiqzvjhgkujbttw",
)

print("SiliconFlow API mode enabled.")


def pil_to_data_url(img: Image.Image) -> str:
    """Convert PIL image to data URL string for API transport."""
    buffered = io.BytesIO()
    img.convert("RGB").save(buffered, format="JPEG", quality=92)
    img_b64 = base64.b64encode(buffered.getvalue()).decode("utf-8")
    return f"data:image/jpeg;base64,{img_b64}"


def extract_json_block(text: str) -> dict:
    """Extract first valid JSON object from model text output."""
    text = text.strip()

    try:
        return json.loads(text)
    except json.JSONDecodeError:
        pass

    fenced = re.search(r"```(?:json)?\s*(\{[\s\S]*?\})\s*```", text)
    if fenced:
        return json.loads(fenced.group(1))

    inline = re.search(r"(\{[\s\S]*\})", text)
    if inline:
        return json.loads(inline.group(1))

    raise ValueError("No valid JSON found in model response")


def call_siliconflow_compare(img1: Image.Image, img2: Image.Image, action_desc: str):
    """Call SiliconFlow API and return (score_img1, score_img2, explanation)."""
    img1_data_url = pil_to_data_url(img1)
    img2_data_url = pil_to_data_url(img2)

    system_prompt = (
        "You are an image-comparison judge. Given two images and an action description, "
        "decide which image is closer to the completed state of the action. "
        "Return only JSON with keys: score_img1, score_img2, explanation. "
        "score_img1 and score_img2 must be numbers in [0,1]."
    )

    payload = {
        "model": MODEL_ID,
        "temperature": 0.1,
        "response_format": {"type": "json_object"},
        "messages": [
            {"role": "system", "content": system_prompt},
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": (
                            "Action description: "
                            f"{action_desc}\n"
                            "Please compare image1 and image2 and output JSON only."
                        ),
                    },
                    {"type": "image_url", "image_url": {"url": img1_data_url}},
                    {"type": "image_url", "image_url": {"url": img2_data_url}},
                ],
            },
        ],
    }

    data = json.dumps(payload).encode("utf-8")
    request = urllib.request.Request(
        API_URL,
        data=data,
        headers={
            "Authorization": f"Bearer {API_KEY}",
            "Content-Type": "application/json",
        },
        method="POST",
    )

    try:
        with urllib.request.urlopen(request, timeout=120) as response:
            raw = response.read().decode("utf-8")
    except urllib.error.HTTPError as e:
        error_body = e.read().decode("utf-8", errors="ignore")
        raise RuntimeError(f"HTTP {e.code}: {error_body}") from e
    except urllib.error.URLError as e:
        raise RuntimeError(f"Network error: {e.reason}") from e

    api_result = json.loads(raw)
    content = api_result["choices"][0]["message"]["content"]
    parsed = extract_json_block(content)

    score_img1 = float(parsed.get("score_img1", 0.0))
    score_img2 = float(parsed.get("score_img2", 0.0))
    explanation = str(parsed.get("explanation", "No explanation provided."))

    score_img1 = max(0.0, min(1.0, score_img1))
    score_img2 = max(0.0, min(1.0, score_img2))

    return score_img1, score_img2, explanation


# ==========================================
# 2. Define Core Inference Function
# ==========================================
def compare_images(img1, img2, action_desc):
    """
    img1: PIL.Image
    img2: PIL.Image
    action_desc: string
    """
    # 1. Input Validation
    if img1 is None or img2 is None:
        return "⚠️ Please ensure both images (img1 and img2) are uploaded!"
    if not action_desc.strip():
        return "⚠️ Please enter an action description!"

    try:
        score_img1, score_img2, api_explanation = call_siliconflow_compare(
            img1=img1,
            img2=img2,
            action_desc=action_desc,
        )

        # 2. Logic Judgment
        if score_img1 > score_img2:
            result = "img1"
            explanation = (
                "The model evaluates that **Image 1 (img1)** is closer to the completed "
                f"state of the described task.\n\nModel rationale: {api_explanation}"
            )
        elif score_img2 > score_img1:
            result = "img2"
            explanation = (
                "The model evaluates that **Image 2 (img2)** is closer to the completed "
                f"state of the described task.\n\nModel rationale: {api_explanation}"
            )
        else:
            result = "Tie"
            explanation = (
                "The model evaluates both images equally regarding the task completion."
                f"\n\nModel rationale: {api_explanation}"
            )

        # 3. Format Output
        return (
            f"### **Output:** {result}\n\n"
            f"**Score(img1):** {score_img1:.3f}  \n"
            f"**Score(img2):** {score_img2:.3f}\n\n"
            f"**Explanation:** {explanation}"
        )
    
    except Exception as e:
        return f"❌ An error occurred during inference: {str(e)}"

# ==========================================
# 3. Build Gradio Interface
# ==========================================
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    # Changed Title to the exact text you requested
    gr.Markdown("# 🧠 SpatialLogic Reasoning Demo")
    # Changed sub-title to English
    gr.Markdown("Upload two images and input an action description. The model will predict which image is closer to the completed state of the task.")
    
    with gr.Row():
        with gr.Column():
            # Changed labels to English
            img1_input = gr.Image(type="pil", label="Image 1 (img1)")
        with gr.Column():
            img2_input = gr.Image(type="pil", label="Image 2 (img2)")
            
    # Changed label and placeholder to English
    text_input = gr.Textbox(
        label="Action Description", 
        placeholder="e.g., Put the apple on the table"
    )
    
    # Changed button text to English
    submit_btn = gr.Button("🤖 Evaluate", variant="primary")
    
    # Output area
    output_text = gr.Markdown(label="Result")

    # Bind function to components
    submit_btn.click(
        fn=compare_images, 
        inputs=[img1_input, img2_input, text_input], 
        outputs=output_text
    )

if __name__ == "__main__":
    demo.launch()