import gradio as gr
import uuid
from datetime import datetime
import pandas as pd
from model_handler import ModelHandler
from tab_chat import create_chat_tab
from tab_code import create_code_tab
from tab_smart_writer import create_smart_writer_tab
from tab_test import run_model_handler_test, run_clear_chat_test

def get_history_df(history):
    if not history:
        return pd.DataFrame({'ID': [], 'Conversation': []})
    df = pd.DataFrame(history)
    return df[['id', 'title']].rename(columns={'id': 'ID', 'title': 'Conversation'})

def on_app_load(history, conv_id):
    """
    Handles the application's initial state on load.
    - If no history exists, creates a new conversation.
    - If the last conversation ID is invalid, loads the most recent one.
    - Otherwise, loads the last active conversation.
    """
    if not history:
        # First time ever loading, create a new chat
        conv_id = str(uuid.uuid4())
        new_convo = { "id": conv_id, "title": "New Conversation", "messages": [], "timestamp": datetime.now().isoformat() }
        history = [new_convo]
        return conv_id, history, gr.update(value=get_history_df(history)), []

    # Check if the last used conv_id is valid
    if conv_id and any(c["id"] == conv_id for c in history):
        # It's valid, load it
        for convo in history:
            if convo["id"] == conv_id:
                return conv_id, history, gr.update(value=get_history_df(history)), convo["messages"]
    
    # Last used conv_id is invalid or doesn't exist, load the most recent conversation
    most_recent_convo = history[0] # Assumes history is sorted by timestamp desc
    conv_id = most_recent_convo["id"]
    return conv_id, history, gr.update(value=get_history_df(history)), most_recent_convo["messages"]

    
CSS = """

#chatbot {
  height: calc(100vh - 21px - 16px);
  max-height: 1500px;
}

footer { 
    display: none !important; 
}
"""

if __name__ == "__main__":
    # Instantiate the model handler with the configuration
    model_handler = ModelHandler()

    with gr.Blocks(theme=gr.themes.Soft(),
                   css=CSS,
                   head="",
                   head_paths=['./static/toastify.html', './static/app.html'],
                   fill_height=True,
                   fill_width=True) as demo:
        with gr.Tabs(elem_id='indicator-space-app') as tabs:

            with gr.TabItem("文本聊天") as chat_tab:
                conversation_store, current_conversation_id, history_df, chatbot = create_chat_tab()

            chat_tab.select(
                fn=None,
                js="() => {window.dispatchEvent(new CustomEvent('tabSelect.chat')); console.log('this'); return null;}",
            )

            with gr.TabItem("代码生成") as code_tab:
                create_code_tab()

            code_tab.select(
                fn=None,
                js="() => {window.dispatchEvent(new CustomEvent('tabSelect.code')); return null;}",
            )

            with gr.TabItem("写作助手") as writer_tab:
                create_smart_writer_tab()

            writer_tab.select(
                fn=None,
                js="() => {window.dispatchEvent(new CustomEvent('tabSelect.writing')); return null;}",
            )

            with gr.TabItem("测试"):
                gr.Markdown("# 功能测试")
                with gr.Column():
                    test_log_output = gr.Textbox(label="测试日志", interactive=False, lines=10)
                    gr.Button("运行 ModelHandler 测试").click(run_model_handler_test, outputs=test_log_output)
                    gr.Button("运行 清除聊天 测试").click(run_clear_chat_test, outputs=test_log_output)

        # Bind on_app_load to demo.load
        demo.load(
            on_app_load,
            inputs=[conversation_store, current_conversation_id],
            outputs=[current_conversation_id, conversation_store, history_df, chatbot],
            js="() => {window.dispatchEvent(new CustomEvent('appStart')); console.log('appStart'); return {};}"
        )

    # Launch the Gradio application
    demo.launch()