import os
import gradio as gr
import asyncio
from typing import Optional
from PIL import Image
import base64
import io
import requests

# --- Configuración de claves ---
class Config:
    def __init__(self):
        self.SAMBANOVA_API_KEY = os.getenv("SAMBANOVA_API_KEY")
        self.BRIA_API_TOKEN = os.getenv("BRIA_API_TOKEN")
        self.validate()

    def validate(self):
        print("🔍 Verificando claves de API...")
        if self.SAMBANOVA_API_KEY:
            print("✅ SAMBANOVA_API_KEY configurada")
        else:
            print("❌ SAMBANOVA_API_KEY no encontrada")
        if self.BRIA_API_TOKEN:
            print("✅ BRIA_API_TOKEN configurado")
        else:
            print("⚠️ BRIA_API_TOKEN no configurado → generación de imágenes desactivada")

config = Config()

# --- Cliente SambaNova (usando SDK oficial) ---
SambaNova = None
sn_client = None

try:
    from sambanova import SambaNova
    if config.SAMBANOVA_API_KEY:
        sn_client = SambaNova(api_key=config.SAMBANOVA_API_KEY)
        print("✅ Cliente SambaNova inicializado")
except Exception as e:
    print(f"❌ Error al inicializar SambaNova: {e}")
    sn_client = None

# --- Herramientas ---
class AI_Tools:
    def __init__(self):
        self.sn_client = sn_client

    def generate_text(self, prompt: str) -> str:
        if not self.sn_client:
            return "❌ SambaNova no disponible. Verifica tu clave API."
        try:
            # ⚠️ Usa un modelo REAL de tu cuenta en SambaNova Cloud
            # Ejemplos comunes: "Meta-Llama-3.1-8B-Instruct", "Llama-3.2-3B-Instruct"
            # Reemplaza "Maverick" por un modelo válido que veas en tu dashboard
            response = self.sn_client.chat.completions.create(
                model="Meta-Llama-3.1-8B-Instruct",  # ✅ CAMBIA ESTO si usas otro modelo
                messages=[{"role": "user", "content": prompt}],
                temperature=0.7,
                max_tokens=500
            )
            return response.choices[0].message.content.strip()
        except Exception as e:
            return f"❌ Error en SambaNova: {e}"

    def generate_image(self, prompt: str) -> Optional[Image.Image]:
        if not config.BRIA_API_TOKEN:
            return None
        try:
            url = "https://api.bria.ai/v1/generate"  # ✅ sin espacio
            headers = {"Authorization": f"Bearer {config.BRIA_API_TOKEN}"}
            json_data = {"prompt": prompt, "options": {"resolution": "512x512"}}
            response = requests.post(url, headers=headers, json=json_data, timeout=30)
            if response.status_code == 200:
                img_b64 = response.json().get("image_base64")
                if img_b64:
                    img_bytes = base64.b64decode(img_b64)
                    return Image.open(io.BytesIO(img_bytes))
        except Exception as e:
            print(f"⚠️ Error en Bria: {e}")
        return None

tools = AI_Tools()

# --- Procesamiento (síncrono para Gradio) ---
def process_input(message: str, history: list) -> list:
    msg_lower = message.lower()
    
    # Generar imagen
    if any(kw in msg_lower for kw in ["imagen", "foto", "genera imagen", "dibuja"]):
        img = tools.generate_image(message)
        if img:
            return history + [(message, ("", img))]
        return history + [(message, "❌ No se pudo generar la imagen. ¿Tienes BRIA_API_TOKEN?")]
    
    # Generar texto con SambaNova
    response = tools.generate_text(message)
    return history + [(message, response)]

# --- Interfaz Gradio ---
with gr.Blocks(title="AI Assistant", theme=gr.themes.Soft()) as app:
    gr.Markdown("# 🤖 AI Assistant (SambaNova + Bria)")
    samba_status = "✅ SambaNova activo" if tools.sn_client else "❌ SambaNova inactivo"
    bria_status = "✅ Bria activo" if config.BRIA_API_TOKEN else "❌ Bria inactivo"
    gr.Markdown(f"**Estado**: {samba_status} | {bria_status}")
    
    chatbot = gr.Chatbot(height=450)
    msg = gr.Textbox(
        label="Escribe tu mensaje",
        placeholder="Ej: 'Explica la relatividad' o 'Genera una imagen de un dragón en Marte'",
        lines=2
    )

    gr.Examples([
        "¿Qué es el aprendizaje automático?",
        "Genera una imagen de un bosque encantado al atardecer",
        "Escribe un breve cuento sobre un robot solitario"
    ], inputs=msg)

    def clear(): return []

    gr.Button("Limpiar chat").click(clear, outputs=chatbot)
    
    # Con queue=True, evita bloqueos
    msg.submit(process_input, [msg, chatbot], chatbot, queue=True)
    gr.Button("Enviar").click(process_input, [msg, chatbot], chatbot, queue=True)

# --- Ejecución ---
if __name__ == "__main__":
    app.queue().launch(server_name="0.0.0.0", server_port=7860)