import gradio as gr
import base64
import io
import requests

def solve_stem_problem(api_key, image, subject):
    if not api_key.strip():
        return "❌ Error: API key is required."

    detectives = {
        "math": "Algebra Ace",
        "physics": "Physics Phantom",
        "chemistry": "Chemistry Clue-finder",
        "coding": "Code Cracker"
    }
    detective = detectives.get(subject, "Algebra Ace")

    try:
        # Convert image to base64 string
        buffered = io.BytesIO()
        image.save(buffered, format="PNG")
        encoded_image = base64.b64encode(buffered.getvalue()).decode('utf-8')
        image_url_data = f"data:image/png;base64,{encoded_image}"
    except Exception as e:
        return f"❌ Error encoding image: {str(e)}"

    try:
        headers = {
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json",
            "HTTP-Referer": "https://stem-sleuth.example.com",
            "X-Title": "STEM Sleuth"
        }

        payload = {
            "model": "google/gemini-2.0-flash-exp:free",
            "messages": [
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "text",
                            "text": f"Act as {detective} and solve this {subject} problem step-by-step with a detective narrative."
                        },
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": image_url_data
                            }
                        }
                    ]
                }
            ]
        }

        response = requests.post("https://openrouter.ai/api/v1/chat/completions", headers=headers, json=payload)

        if response.status_code != 200:
            return f"❌ API Error {response.status_code}: {response.text}"

        completion = response.json()
        message = completion.get("choices", [{}])[0].get("message", {}).get("content")

        if not message:
            return "⚠️ No response content received. Try another image or check the model output."

        return message

    except Exception as e:
        return f"❌ API call failed: {str(e)}"

# Launch Gradio Interface
iface = gr.Interface(
    fn=solve_stem_problem,
    inputs=[
        gr.Textbox(label="🔐 OpenRouter API Key", type="password"),
        gr.Image(label="🖼️ Upload STEM Problem Image", type="pil"),
        gr.Dropdown(["math", "physics", "chemistry", "coding"], label="📚 Select Subject")
    ],
    outputs=gr.Textbox(label="🕵️‍♂️ Detective's Solution"),
    title="🧠 STEM Sleuth Solver",
    description="Upload a math, physics, chemistry, or coding problem image and solve it with a detective twist using OpenRouter Gemini model."
)

iface.launch()