import gradio as gr
import requests
import json
import os
import base64

# Set your OpenRouter API key in the Space's secrets as "OPENROUTER_API_KEY"
OPENROUTER_API_KEY = os.getenv("OPENROUTER_API_KEY")
HEADERS = {
    "Authorization": f"Bearer {OPENROUTER_API_KEY}",
    "HTTP-Referer": "https://huggingface.co/spaces/YOUR_SPACE",  # Optional
    "X-Title": "CrispChat"  # Optional
}

# List of free OpenRouter models
FREE_MODELS = {
    "Google: Gemini Pro 2.5 Experimental (free)": ("google/gemini-2.5-pro-exp-03-25:free", 1000000),
    "DeepSeek: DeepSeek V3 (free)": ("deepseek/deepseek-chat:free", 131072),
    "Meta: Llama 3.2 11B Vision Instruct (free)": ("meta-llama/llama-3.2-11b-vision-instruct:free", 131072),
    "Qwen: Qwen2.5 VL 72B Instruct (free)": ("qwen/qwen2.5-vl-72b-instruct:free", 131072),
}

def query_openrouter_model(model_id, prompt, image=None):
    # If there's an image, embed it in the prompt as text
    if image is not None:
        with open(image, "rb") as f:
            b64_img = base64.b64encode(f.read()).decode("utf-8")
        prompt += f"\n[User attached an image: data:image/png;base64,{b64_img}]"

    messages = [
        {"role": "user", "content": prompt}
    ]

    payload = {
        "model": model_id,
        "messages": messages
    }

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

    try:
        response.raise_for_status()
        data = response.json()
        return data["choices"][0]["message"]["content"]
    except Exception as e:
        return f"Error: {str(e)}\n{response.text}"

def chat_interface(prompt, image, model_label):
    model_id, _ = FREE_MODELS[model_label]
    return query_openrouter_model(model_id, prompt, image)

with gr.Blocks(title="CrispChat") as demo:
    gr.Markdown("""
    # 🌟 CrispChat
    Multi-modal chat with free OpenRouter models
    """)

    with gr.Row():
        prompt = gr.Textbox(label="Enter your message", lines=4, placeholder="Ask me anything...")
        image = gr.Image(type="filepath", label="Optional image input")

    model_choice = gr.Dropdown(
        choices=list(FREE_MODELS.keys()),
        value="Google: Gemini Pro 2.5 Experimental (free)",
        label="Select model"
    )

    output = gr.Textbox(label="Response", lines=6)

    submit = gr.Button("Submit")
    submit.click(fn=chat_interface, inputs=[prompt, image, model_choice], outputs=output)

if __name__ == "__main__":
    # Hide the API documentation to avoid schema issues
    demo.launch(show_api=False)