ui_module.py

import gradio as gr
import uuid
import logging
import re
import json
import os
import tempfile
from agent_module import chat_agent_stream
from translation_module import localize_init

logger = logging.getLogger("app.ui")

PREMIUM_CSS = """
body, .gradio-container { background-color: #0b0f19 !important; color: #e0e6ed !important; font-family: 'Inter', sans-serif !important; }
.main-container { max-width: 900px !important; margin: 0 auto !important; min-height: 600px !important; max-height: 90vh !important; display: flex !important; flex-direction: column !important; padding: 10px !important; background: #0b0f19 !important; }
.header-row { display: flex !important; align-items: center !important; justify-content: space-between !important; gap: 10px !important; margin-bottom: 10px !important; padding: 10px !important; background: rgba(255, 255, 255, 0.05) !important; border-radius: 12px !important; }
.chat-window { flex-grow: 1 !important; border-radius: 12px !important; border: 1px solid rgba(255, 255, 255, 0.1) !important; background: rgba(0, 0, 0, 0.2) !important; overflow-y: auto !important; }
.input-area { margin-top: 10px !important; background: rgba(255, 255, 255, 0.05) !important; padding: 10px !important; border-radius: 12px !important; flex-shrink: 0 !important; }
footer { display: none !important; }
@media (max-width: 768px) { .main-container { padding: 5px !important; height: auto !important; max-height: none !important; } .header-row { padding: 5px !important; flex-wrap: wrap !important; } }
"""

def save_and_clear(msg):
    return "", msg

def export_chat(history):
    if not history: return None
    with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False) as tf:
        json.dump(history, tf, indent=2)
        return tf.name

def import_chat(file):
    if file is None: return gr.update(), []
    try:
        with open(file.name, 'r') as f:
            data = json.load(f)
            if isinstance(data, list):
                return data, data
    except Exception as e:
        logger.error(f"Import error: {e}")
    return gr.update(), gr.update()

def clear_messages(lang_name="English"):
    name_map = {"German": "de", "French": "fr", "Spanish": "es", "Italian": "it", 
                "Portuguese": "pt", "Russian": "ru", "Japanese": "ja", "Chinese": "zh"}
    code = name_map.get(lang_name, "en")
    loc = localize_init(code)
    welcome = [{"role": "assistant", "content": loc["welcome"]}]
    return welcome, welcome

def chat_wrapper(message, history, short_answers=False, lang="English"):
    if history is None: history = []
    
    history.append({"role": "user", "content": message})
    history.append({"role": "assistant", "content": ""})

    # Intent detection is now handled internally by chat_agent_stream
    for part in chat_agent_stream(message, history[:-2], user_lang=lang, short_answers=short_answers):
        if part == "__TURN_END__":
            history.append({"role": "assistant", "content": ""})
            yield history, history
        else:
            history[-1]["content"] = part
            yield history, history
    
    yield history, history

def on_load(request: gr.Request):
    headers = dict(request.headers)
    lang_code = "en"
    try:
        if "accept-language" in headers:
            lang_code = headers.get("accept-language").split(",")[0].split("-")[0]
    except: pass

    loc = localize_init(lang_code)
    welcome = [{"role": "assistant", "content": loc["welcome"]}]
    return welcome, loc["lang"], gr.update(label=loc["label_short"]), gr.update(placeholder=loc["placeholder"])

def build_demo(enable_reload=True):
    with gr.Blocks(css=PREMIUM_CSS, title="Sage 6.5 - Oracle Intermediary") as demo:
        history_state = gr.State([])
        lang_state = gr.State("English")
        saved_msg = gr.State("")

        # Default Static Localization for Immediate Render
        default_loc = localize_init("en") 
        default_welcome = [{"role": "assistant", "content": default_loc["welcome"]}]

        with gr.Column(elem_classes="main-container"):
            with gr.Row(elem_classes="header-row"):
                gr.Markdown("## 🔮 Sage 6.5")
                with gr.Row():
                    import_btn = gr.UploadButton("📥 Import", file_types=[".json"], size="sm")
                    export_btn = gr.DownloadButton("📤 Export", size="sm")
            
            # BEST PRACTICE: Init with value, don't wait for load
            chatbot = gr.Chatbot(value=default_welcome, elem_classes="chat-window", type="messages", bubble_full_width=False, show_label=False)
            
            with gr.Row(elem_classes="input-area"):
                msg_input = gr.Textbox(placeholder=default_loc["placeholder"], show_label=False, container=False, scale=4)
                submit_btn = gr.Button("➤", scale=1, variant="primary")
            
            short_answers = gr.Checkbox(label=default_loc["label_short"], value=False)

        # Wire Events
        msg_input.submit(save_and_clear, [msg_input], [msg_input, saved_msg], queue=False).then(
            chat_wrapper, [saved_msg, history_state, short_answers, lang_state], [chatbot, history_state]
        )
        submit_btn.click(save_and_clear, [msg_input], [msg_input, saved_msg], queue=False).then(
            chat_wrapper, [saved_msg, history_state, short_answers, lang_state], [chatbot, history_state]
        )

        export_btn.click(export_chat, [history_state], [export_btn])
        import_btn.upload(import_chat, [import_btn], [chatbot, history_state])
        chatbot.clear(clear_messages, [lang_state], [chatbot, history_state])

        # State Sync must happen after init
        # We manually sync state with the default value so backend has it
        def sync_init():
             return default_welcome, default_welcome
        
        # If reload is enabled (client-side), we still try to localize dynamically, but we have a safe fallback now.
        if enable_reload:
            demo.load(on_load, None, [chatbot, lang_state, short_answers, msg_input])
        else:
            # If no reload, ensure history state matches visual init
            demo.load(sync_init, None, [history_state, chatbot]) # Sync state

        api_history = gr.Chatbot(visible=False, type="messages")
        api_chat_btn = gr.Button("API CHAT", visible=False)
        api_chat_btn.click(chat_wrapper, [msg_input, api_history, short_answers, lang_state], [chatbot, history_state], api_name="chat")
    
    return demo