import os
import base64
import gradio as gr
import requests
import json
from io import BytesIO
from PIL import Image

# Get API key from environment variable for security
OPENROUTER_API_KEY = os.environ.get("OPENROUTER_API_KEY", "")

# Simplified model list
models = [
    ("Google Gemini Pro 2.0", "google/gemini-2.0-pro-exp-02-05:free"),
    ("Google Gemini 2.5 Pro", "google/gemini-2.5-pro-exp-03-25:free"),
    ("Meta Llama 3.2 Vision", "meta-llama/llama-3.2-11b-vision-instruct:free"),
    ("Qwen 2.5 VL", "qwen/qwen2.5-vl-72b-instruct:free"),
    ("DeepSeek R1", "deepseek/deepseek-r1:free"),
    ("Mistral 3.1", "mistralai/mistral-small-3.1-24b-instruct:free")
]

def get_response(message, history, model_name, image=None, file=None):
    """Simple function to get response from API"""
    # Find model ID from name
    model_id = next((mid for name, mid in models if name == model_name), models[0][1])
    
    # Format messages from history
    messages = []
    for human, ai in history:
        messages.append({"role": "user", "content": human})
        if ai:  # Only add if there's a response
            messages.append({"role": "assistant", "content": ai})
    
    # Process file if provided
    if file:
        try:
            with open(file.name, 'r', encoding='utf-8') as f:
                file_content = f.read()
                message = f"{message}\n\nFile content:\n```\n{file_content}\n```"
        except Exception as e:
            message = f"{message}\n\nError reading file: {str(e)}"
    
    # Process image if provided
    if image is not None:
        try:
            # Convert image to base64
            buffered = BytesIO()
            image.save(buffered, format="JPEG")
            base64_image = base64.b64encode(buffered.getvalue()).decode("utf-8")
            
            # Create multimodal content
            content = [
                {"type": "text", "text": message},
                {
                    "type": "image_url",
                    "image_url": {
                        "url": f"data:image/jpeg;base64,{base64_image}"
                    }
                }
            ]
            messages.append({"role": "user", "content": content})
        except Exception as e:
            messages.append({"role": "user", "content": f"{message}\n\nError processing image: {str(e)}"})
    else:
        messages.append({"role": "user", "content": message})
    
    # Make API call (non-streaming for reliability)
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {OPENROUTER_API_KEY}",
        "HTTP-Referer": "https://huggingface.co/spaces",
    }
    
    data = {
        "model": model_id,
        "messages": messages,
        "temperature": 0.7,
        "max_tokens": 1000
    }
    
    try:
        response = requests.post(
            "https://openrouter.ai/api/v1/chat/completions",
            headers=headers,
            json=data,
            timeout=60
        )
        response.raise_for_status()
        
        result = response.json()
        reply = result.get("choices", [{}])[0].get("message", {}).get("content", "No response")
        
        return reply
    except Exception as e:
        return f"Error: {str(e)}"

# Create ultra simple interface
with gr.Blocks() as demo:
    gr.Markdown("# 🔆 CrispChat")
    
    chatbot = gr.Chatbot(height=450)
    
    with gr.Row():
        with gr.Column(scale=3):
            msg = gr.Textbox(
                placeholder="Type your message here...",
                lines=3,
                label="Message"
            )
        
        with gr.Column(scale=1):
            model = gr.Dropdown(
                choices=[name for name, _ in models],
                value=models[0][0],
                label="Model"
            )
    
    with gr.Row():
        with gr.Column(scale=1):
            img = gr.Image(type="pil", label="Image (optional)")
        
        with gr.Column(scale=1):
            file = gr.File(label="Text File (optional)")
    
    with gr.Row():
        submit = gr.Button("Send")
        clear = gr.Button("Clear")
    
    # Events
    submit.click(
        fn=get_response,
        inputs=[msg, chatbot, model, img, file],
        outputs=chatbot
    ).then(
        lambda: "", None, None,
        outputs=[msg, img, file]
    )
    
    msg.submit(
        fn=get_response,
        inputs=[msg, chatbot, model, img, file],
        outputs=chatbot
    ).then(
        lambda: "", None, None,
        outputs=[msg, img, file]
    )
    
    clear.click(lambda: [], outputs=chatbot)

# Define FastAPI endpoint
from fastapi import FastAPI
from pydantic import BaseModel

app = FastAPI()

class GenerateRequest(BaseModel):
    message: str
    model: str = None
    image_data: str = None

@app.post("/api/generate")
async def api_generate(request: GenerateRequest):
    """Simple API endpoint"""
    model_id = request.model or models[0][1]
    
    messages = []
    
    # Process image if provided
    if request.image_data:
        try:
            # Decode base64 image
            image_bytes = base64.b64decode(request.image_data)
            image = Image.open(BytesIO(image_bytes))
            
            # Re-encode to ensure proper format
            buffered = BytesIO()
            image.save(buffered, format="JPEG")
            base64_image = base64.b64encode(buffered.getvalue()).decode("utf-8")
            
            content = [
                {"type": "text", "text": request.message},
                {
                    "type": "image_url",
                    "image_url": {
                        "url": f"data:image/jpeg;base64,{base64_image}"
                    }
                }
            ]
            messages.append({"role": "user", "content": content})
        except Exception as e:
            return {"error": f"Image processing error: {str(e)}"}
    else:
        messages.append({"role": "user", "content": request.message})
    
    # Make API call
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {OPENROUTER_API_KEY}",
        "HTTP-Referer": "https://huggingface.co/spaces",
    }
    
    data = {
        "model": model_id,
        "messages": messages,
        "temperature": 0.7
    }
    
    try:
        response = requests.post(
            "https://openrouter.ai/api/v1/chat/completions",
            headers=headers,
            json=data,
            timeout=60
        )
        response.raise_for_status()
        
        result = response.json()
        reply = result.get("choices", [{}])[0].get("message", {}).get("content", "No response")
        
        return {"response": reply}
    except Exception as e:
        return {"error": f"Error: {str(e)}"}

# Mount Gradio app
app = gr.mount_gradio_app(app, demo, path="/")

# Launch the app
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)