import streamlit as st
import os
import requests
from io import BytesIO
from openai import OpenAI

XAI_API_KEY = os.getenv("XAI_API_KEY", "")

# El client se inicializa solo si hay key
def get_client():
    key = XAI_API_KEY or st.session_state.get("api_key", "")
    if not key:
        return None
    return OpenAI(api_key=key, base_url="https://api.x.ai/v1")

def get_key():
    return XAI_API_KEY or st.session_state.get("api_key", "")

st.set_page_config(page_title="Grok Voice Studio", page_icon="🔊", layout="wide")
st.title("🔊 Grok Voice Studio - TTS, Chat y Clonador xAI")
st.markdown("**Plataforma completa con voz nativa Grok, chat y clonación de voz**")

# Si no hay key en env, pedir al usuario
if not XAI_API_KEY:
    with st.sidebar:
        st.subheader("🔑 Configuración")
        user_key = st.text_input("XAI_API_KEY", type="password", placeholder="xai-...")
        if user_key:
            st.session_state["api_key"] = user_key
            st.success("Key configurada")

def call_tts(text: str, voice: str = "eve", language: str = "es", speed: float = 1.0):
    """
    xAI TTS — intentamos el endpoint oficial.
    Si no existe todavía (404/501), devolvemos None y mostramos aviso.
    """
    key = get_key()
    if not key:
        return None, "Sin API key"
    
    try:
        # Intentar con el endpoint nativo de xAI
        r = requests.post(
            "https://api.x.ai/v1/audio/speech",
            headers={"Authorization": f"Bearer {key}", "Content-Type": "application/json"},
            json={
                "model": "grok-tts",          # nombre tentativo según docs xAI
                "input": text,
                "voice": voice,
                "response_format": "mp3",
                "speed": speed
            },
            timeout=30
        )
        if r.status_code == 200:
            return BytesIO(r.content), None
        
        # Fallback: endpoint legacy que usaba el app original
        r2 = requests.post(
            "https://api.x.ai/v1/tts",
            headers={"Authorization": f"Bearer {key}", "Content-Type": "application/json"},
            json={"text": text, "voice_id": voice, "language": language, "speed": speed},
            timeout=30
        )
        if r2.status_code == 200:
            return BytesIO(r2.content), None
        
        # Si ninguno funciona, devolver el error real
        return None, f"TTS no disponible aún (HTTP {r.status_code}): {r.text[:300]}"
    except Exception as e:
        return None, str(e)


tab1, tab2, tab3, tab4 = st.tabs(["📝 TTS Directo", "💬 Chat + Voz", "🗣️ Clonador", "ℹ️ Info API"])

# ─── TAB 1: TTS DIRECTO ─────────────────────────────────────────────────────
with tab1:
    st.subheader("Lector Directo TTS xAI")
    texto = st.text_area("Texto:", height=180, placeholder="Escribe el texto que quieres escuchar...")
    
    c1, c2, c3 = st.columns(3)
    with c1:
        voice = st.selectbox("Voz", ["eve", "ara", "rex", "sal", "leo"], index=0)
    with c2:
        lang = st.selectbox("Idioma", ["es", "en", "fr", "pt", "it", "auto"], index=0)
    with c3:
        speed = st.slider("Velocidad", 0.5, 2.0, 1.0, 0.1)

    if st.button("🎙️ Generar Audio", type="primary", key="btn_tts"):
        if not get_key():
            st.error("Configura tu XAI_API_KEY primero (sidebar o variable de entorno)")
        elif texto.strip():
            with st.spinner("Procesando con Grok TTS..."):
                audio_bytes, err = call_tts(texto, voice, lang, speed)
                if audio_bytes:
                    st.audio(audio_bytes, format="audio/mp3")
                    st.success("✅ Audio generado")
                else:
                    st.error(f"Error TTS: {err}")
                    st.info("💡 El endpoint TTS de xAI puede no estar disponible en tu plan actual. Consulta https://docs.x.ai")
        else:
            st.warning("Escribe algún texto primero")

# ─── TAB 2: CHAT + VOZ ───────────────────────────────────────────────────────
with tab2:
    st.subheader("Chat Grok con Voz")
    
    # Selector de modelo — grok-4 puede no estar disponible en todos los planes
    model_name = st.selectbox("Modelo", ["grok-3", "grok-3-mini", "grok-2-1212"], index=0, key="chat_model")
    tts_on = st.toggle("🔊 Leer respuestas en voz alta", value=True)
    
    if "msgs" not in st.session_state:
        st.session_state.msgs = []
    
    # Botón limpiar chat
    if st.button("🗑️ Limpiar chat", key="clear_chat"):
        st.session_state.msgs = []
        st.rerun()
    
    for m in st.session_state.msgs:
        with st.chat_message(m["role"]):
            st.write(m["content"])
    
    prompt = st.chat_input("Mensaje...")
    if prompt:
        if not get_key():
            st.error("Configura tu XAI_API_KEY primero")
        else:
            st.session_state.msgs.append({"role": "user", "content": prompt})
            with st.chat_message("user"):
                st.write(prompt)
            
            with st.chat_message("assistant"):
                with st.spinner("Pensando..."):
                    try:
                        c = get_client()
                        resp = c.chat.completions.create(
                            model=model_name,
                            messages=[{"role": m["role"], "content": m["content"]} for m in st.session_state.msgs],
                            temperature=0.7,
                            max_tokens=2048
                        )
                        txt = resp.choices[0].message.content
                        st.write(txt)
                        st.session_state.msgs.append({"role": "assistant", "content": txt})
                        
                        if tts_on:
                            with st.spinner("Generando voz..."):
                                audio_bytes, err = call_tts(txt[:4000])  # límite seguro
                                if audio_bytes:
                                    st.audio(audio_bytes, format="audio/mp3")
                                else:
                                    st.caption(f"⚠️ TTS no disponible: {err}")
                    except Exception as e:
                        st.error(f"Error chat: {e}")

# ─── TAB 3: CLONADOR ─────────────────────────────────────────────────────────
with tab3:
    st.subheader("Clonador de Voz xAI")
    st.info("ℹ️ La clonación de voz requiere acceso al endpoint `/v1/custom-voices` de xAI, actualmente en beta privada.")
    
    aud = st.file_uploader("Audio referencia (wav/mp3, mín. 30s recomendado)", type=["wav", "mp3"])
    vname = st.text_input("Nombre para la voz", "MiVoz")
    
    if aud and st.button("🧬 Clonar Voz", key="btn_clone"):
        if not get_key():
            st.error("Configura tu XAI_API_KEY primero")
        else:
            with st.spinner("Clonando voz... (puede tardar varios segundos)"):
                try:
                    files = {"file": (aud.name, aud.getvalue(), "audio/wav")}
                    cr = requests.post(
                        "https://api.x.ai/v1/custom-voices",
                        headers={"Authorization": f"Bearer {get_key()}"},
                        files=files,
                        data={"name": vname, "language": "es"},
                        timeout=60
                    )
                    if cr.status_code == 200:
                        vid = cr.json().get("voice_id") or cr.json().get("id")
                        st.success(f"✅ Voz creada! ID: `{vid}`")
                        st.session_state["cloned_voice_id"] = vid
                    else:
                        st.error(f"HTTP {cr.status_code}: {cr.text[:400]}")
                        st.info("El endpoint de clonación puede no estar disponible públicamente aún.")
                except Exception as e:
                    st.error(str(e))
    
    # Probar voz clonada si existe en session
    if "cloned_voice_id" in st.session_state:
        st.divider()
        st.write(f"**Voz activa:** `{st.session_state['cloned_voice_id']}`")
        test_text = st.text_area("Texto de prueba", "Hola, esta es mi voz clonada con Grok.")
        if st.button("▶️ Probar voz clonada"):
            with st.spinner("Sintetizando..."):
                audio_bytes, err = call_tts(test_text, voice=st.session_state["cloned_voice_id"])
                if audio_bytes:
                    st.audio(audio_bytes, format="audio/mp3")
                else:
                    st.error(err)

# ─── TAB 4: INFO API ─────────────────────────────────────────────────────────
with tab4:
    st.subheader("Estado de la API xAI")
    
    if st.button("🔍 Verificar conexión API", key="btn_check"):
        if not get_key():
            st.error("Configura tu XAI_API_KEY primero")
        else:
            with st.spinner("Verificando..."):
                try:
                    c = get_client()
                    # Llamada mínima para verificar autenticación
                    models = c.models.list()
                    model_names = [m.id for m in models.data]
                    st.success("✅ API key válida")
                    st.write("**Modelos disponibles en tu cuenta:**")
                    for mn in model_names:
                        st.code(mn)
                except Exception as e:
                    st.error(f"Error: {e}")
    
    st.divider()
    st.markdown("""
    ### Endpoints xAI relevantes
    | Endpoint | Estado |
    |----------|--------|
    | `/v1/chat/completions` | ✅ Disponible |
    | `/v1/models` | ✅ Disponible |
    | `/v1/audio/speech` | ⚠️ Beta / verificar |
    | `/v1/tts` | ⚠️ Beta / verificar |
    | `/v1/custom-voices` | 🔒 Beta privada |
    
    **Docs:** https://docs.x.ai/api  
    **Console:** https://console.x.ai
    """)

st.caption("Grok Voice Studio • Powered by xAI API")