import gradio as gr
from gradio_client import Client, handle_file 
import os

hf_token = os.environ.get("HF_TOKEN")

def safe_client(space_id):
    try:
        return Client(space_id, token=hf_token)
    except Exception as e:
        print(f"⚠️ Connection failed for {space_id}: {e}")
        return None

# Sub-space clients
speech_emo_client = safe_client("E-motionAssistant/ser-wav2vec")
text_emo_client   = safe_client("E-motionAssistant/Space4")
llm_english       = safe_client("E-motionAssistant/TherapyEnglish")
tts_client         = safe_client("E-motionAssistant/Space3")

def main_orchestrator(audio_input, text_input, history):
    if history is None:
        history = []

    # 1. EMOTION LOGIC
    emotion = "Neutral"
    try:
        if audio_input:
            emotion = speech_emo_client.predict(handle_file(audio_input), api_name="/predict")
        elif text_input:
            emotion = text_emo_client.predict(text_input, api_name="/predict")
    except:
        emotion = "Neutral"

    # 2. FORMAT FOR SUB-SPACE (STRICT DICTIONARY)
    # Even if our UI is old, our sub-space (TherapyEnglish) needs the new format
    api_history = []
    for u, b in history:
        api_history.append({"role": "user", "content": [{"type": "text", "text": str(u)}]})
        api_history.append({"role": "assistant", "content": [{"type": "text", "text": str(b)}]})

    bundled_text = f"Context: User is {emotion}. Message: {text_input}"
    current_message = {"role": "user", "content": [{"type": "text", "text": bundled_text}]}

    # 3. CALL LLM
    try:
        response = llm_english.predict(
            message=current_message,
            history=api_history,
            api_name="/chat"
        )
    except Exception as e:
        response = f"LLM Error: {str(e)}"

    # 4. TTS LOGIC
    audio_res = None
    try:
        if tts_client and response:
            audio_res = tts_client.predict(str(response), api_name="/predict")
    except:
        audio_res = None

    # 5. UPDATE HISTORY (Simple list of lists for old Gradio UI)
    history.append([text_input, response])
    return history, audio_res

with gr.Blocks() as demo:
    state = gr.State([])
    with gr.Row():
        with gr.Column():
            audio_in = gr.Audio(label="Voice", type="filepath")
            text_in = gr.Textbox(label="Message")
            btn = gr.Button("Send")
        with gr.Column():
            # REMOVED type="messages" to stop the crash
            chatbot_ui = gr.Chatbot(label="Therapy History") 
            audio_out = gr.Audio(autoplay=True)

    btn.click(main_orchestrator, [audio_in, text_in, state], [chatbot_ui, audio_out])

demo.launch()