import os
import gradio as gr
import base64
from openai import OpenAI

# Configuración inicial
api_key = os.getenv("SAMBANOVA_API_KEY")
client = OpenAI(api_key=api_key, base_url="https://api.sambanova.ai/v1")

# Lista de modelos disponibles
MODELS = {
    "general_fast": "Meta-Llama-3.1-8B-Instruct",
    "general_smart": "Meta-Llama-3.3-70B-Instruct",
    "coding_expert": "DeepSeek-V3.1",
    "coding_alt": "DeepSeek-V3-0324",
    "massive_brain": "gpt-oss-120b",
    "specialized_1": "DeepSeek-V3.1-Terminus",
    "specialized_2": "Llama-3.3-Swallow-70B-Instruct-v0.4",
    "multilingual": "Qwen3-32B",
    "arabic_special": "ALLaM-7B-Instruct-preview",
    "vision_expert": "Llama-4-Maverick-17B-128E-Instruct"
}

# Función auxiliar para imágenes
def encode_image(image_path):
    try:
        with open(image_path, "rb") as image_file:
            return base64.b64encode(image_file.read()).decode('utf-8')
    except Exception as e:
        print(f"Error al codificar la imagen: {e}")
        return None

# Función para seleccionar el modelo más adecuado
def select_best_model(user_prompt):
    try:
        router_system_prompt = (
            "You are an AI router. Analyze the user's request and select the best ID from this list:\n"
            f"{list(MODELS.values())}\n"
            "Rules:\n"
            "- For coding/math: Choose 'DeepSeek-V3.1' or 'DeepSeek-V3-0324'.\n"
            "- For complex logic: Choose 'Meta-Llama-3.3-70B-Instruct' or 'gpt-oss-120b'.\n"
            "- For Japanese/Cultural: Choose 'Llama-3.3-Swallow-70B-Instruct-v0.4'.\n"
            "- For Arabic/Middle East: Choose 'ALLaM-7B-Instruct-preview'.\n"
            "- For General/Fast chat: Choose 'Meta-Llama-3.1-8B-Instruct'.\n"
            "- For Multilingual/Chinese: Choose 'Qwen3-32B'.\n"
            "Reply ONLY with the exact model name strings from the list. Nothing else."
        )

        response = client.chat.completions.create(
            model="Meta-Llama-3.1-8B-Instruct",
            messages=[
                {"role": "system", "content": router_system_prompt},
                {"role": "user", "content": f"User request: {user_prompt}"}
            ],
            temperature=0.0,
            max_tokens=50
        )
        selected_model = response.choices[0].message.content.strip()
        selected_model = selected_model.replace('"', '').replace("'", "")

        if selected_model not in MODELS.values():
            return "Meta-Llama-3.3-70B-Instruct"
        return selected_model
    except Exception as e:
        print(f"Error en router: {e}")
        return "Meta-Llama-3.3-70B-Instruct"

# Lógica del chat
def chat_logic(message, history):
    user_text = message["text"]
    files = message.get("files")

    if not user_text and not files:
        return "Por favor, ingresa un texto o sube una imagen."

    target_model = ""
    messages_payload = []

    if files and len(files) > 0:
        target_model = MODELS["vision_expert"]
        try:
            image_b64 = encode_image(files[0])
            messages_payload = [
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": user_text or "Describe this image."},
                        {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{image_b64}"}}
                    ]
                }
            ]
            yield f"👁️ **Modelo Auto-Seleccionado:** `{target_model}` (Detecté una imagen, mi rey)\n\n"
        except Exception as img_err:
            yield f"¡Ups! No pude procesar la imagen: {img_err}"
            return

    else:
        target_model = select_best_model(user_text)
        messages_payload = [
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": user_text}
        ]
        yield f"🧠 **Modelo Auto-Seleccionado:** `{target_model}` (Analizando tu petición...)\n\n"

    try:
        stream = client.chat.completions.create(
            model=target_model,
            messages=messages_payload,
            temperature=0.1,
            top_p=0.1,
            stream=True
        )

        partial_response = f"🧠 **Modelo Auto-Seleccionado:** `{target_model}`\n\n"

        for chunk in stream:
            content = chunk.choices[0].delta.content
            if content:
                partial_response += content
                yield partial_response

    except Exception as e:
        yield f"¡Chale! Hubo un error en el sistema, carnal: {str(e)}"

# Interfaz Gradio
custom_css = """
:root {
    --primary: #ff6b6b;
    --secondary: #4ecdc4;
    --dark: #1a1a2e;
    --light: #f7f7f7;
}

.gradio-container {
    background: linear-gradient(135deg, var(--dark) 0%, #16213e 50%, #0f3460 100%);
    min-height: 100vh;
    font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
}

.dark .chatbot {
    background: rgba(255, 255, 255, 0.95) !important;
    border-radius: 15px !important;
    box-shadow: 0 8px 32px rgba(0, 0, 0, 0.3) !important;
    backdrop-filter: blur(10px);
    border: 1px solid rgba(255, 255, 255, 0.2);
}

h1 {
    color: white;
    text-align: center;
    text-shadow: 2px 2px 4px rgba(0,0,0,0.5);
    font-size: 2.5em !important;
    margin-bottom: 10px !important;
    background: linear-gradient(45deg, var(--primary), var(--secondary));
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    background-clip: text;
}

h3 {
    color: rgba(255, 255, 255, 0.9) !important;
    text-align: center;
    font-weight: 300 !important;
    margin-top: 0 !important;
}

.textbox {
    border-radius: 10px !important;
    border: 2px solid var(--secondary) !important;
}

button {
    background: linear-gradient(45deg, var(--primary), var(--secondary)) !important;
    border: none !important;
    border-radius: 10px !important;
    color: white !important;
    font-weight: bold !important;
}

button:hover {
    transform: translateY(-2px);
    box-shadow: 0 4px 15px rgba(0, 0, 0, 0.2) !important;
}
"""

with gr.Blocks() as demo:
    gr.Markdown("# 🤖 BATUTO-ART: Selector de IA Inteligente")
    gr.Markdown("### Escribe lo que quieras o sube una foto. Yo elijo la IA perfecta pa' ti.")

    chat = gr.ChatInterface(
        fn=chat_logic,
        multimodal=True,
        textbox=gr.MultimodalTextbox(
            file_types=["image"],
            placeholder="Escribe aquí o sube una foto, papacito...",
        ),
    )

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Soft(), css=custom_css)