"""
LifeFlow AI - Main Application (Fixed Output Mismatch)
✅ 修復 Step 4 回傳值數量不匹配導致的 AttributeError
✅ 確保 outputs 清單包含所有 7 個元件
"""

import gradio as gr
import inspect
from datetime import datetime
from ui.theme import get_enhanced_css
from ui.components.header import create_header
from ui.components.progress_stepper import create_progress_stepper, update_stepper
from ui.components.input_form import create_input_form, toggle_location_inputs
from ui.components.modals import create_settings_modal, create_doc_modal
from ui.renderers import (
    create_agent_dashboard,
    create_summary_card,
)
from core.session import UserSession
from services.planner_service import PlannerService
from services.validator import APIValidator # 記得 import
from config import MODEL_OPTIONS

class LifeFlowAI:
    def __init__(self):
        self.service = PlannerService()

    def cancel_wrapper(self, session_data):
        session = UserSession.from_dict(session_data)
        if session.session_id:
            self.service.cancel_session(session.session_id)
        # 不需要回傳任何東西，這只是一個副作用 (Side Effect) 函數

    def _get_dashboard_html(self, active_agent: str = None, status: str = "idle", message: str = "Waiting") -> str:
        agents = ['team', 'scout', 'optimizer', 'navigator', 'weatherman', 'presenter']
        status_dict = {a: {'status': 'idle', 'message': 'Standby'} for a in agents}
        if active_agent in status_dict:
            status_dict[active_agent] = {'status': status, 'message': message}
            if active_agent != 'team' and status == 'working':
                status_dict['team'] = {'status': 'working', 'message': f'Monitoring {active_agent}...'}
        return create_agent_dashboard(status_dict)

    def _check_api_status(self, session_data):
        session = UserSession.from_dict(session_data)
        has_model_api = bool(session.custom_settings.get("model_api"))
        has_google = bool(session.custom_settings.get("google_maps_api_key"))
        msg = "✅ System Ready" if (has_model_api and has_google) else "⚠️ Missing API Keys"
        return msg

    def _get_gradio_chat_history(self, session):
        history = []
        for msg in session.chat_history:
            history.append({"role": msg["role"], "content": msg["message"]})
        return history

    # ================= Event Wrappers =================

    def analyze_wrapper(self, user_input, auto_loc, lat, lon, session_data):

        session = UserSession.from_dict(session_data)
        iterator = self.service.run_step1_analysis(user_input, auto_loc, lat, lon, session)
        for event in iterator:
            evt_type = event.get("type")
            current_session = event.get("session", session)
            if evt_type == "error":
                gr.Warning(event.get('message'))
                yield (
                    gr.update(visible=True),  # Step 1 Container: 保持顯示
                    gr.update(visible=False),  # Step 2 Container: 保持隱藏
                    gr.update(visible=False),  # Step 3 Container
                    gr.HTML(f"<div style='color:red'>{event.get('message')}</div>"),  # s1_stream
                    gr.update(),  # task_list
                    gr.update(),  # task_summary
                    gr.update(),  # chatbot
                    current_session.to_dict()  # state
                )
                return

            if evt_type == "stream":
                yield (
                    gr.update(visible=True),
                    gr.update(visible=False),
                    gr.update(visible=False),
                    event.get("stream_text", ""),
                    gr.update(),
                    gr.update(),
                    gr.update(),
                    current_session.to_dict()
                )

            elif evt_type == "complete":
                tasks_html = self.service.generate_task_list_html(current_session)

                date_str = event.get("start_time", "N/A")
                high_priority = event.get("high_priority", "N/A")
                total_time = event.get("total_time", "N/A")
                loc_str = event.get("start_location", "N/A")

                summary_html = create_summary_card(len(current_session.task_list),
                                                   high_priority,
                                                   total_time,
                                                   location=loc_str, date=date_str)

                chat_history = self._get_gradio_chat_history(current_session)
                if not chat_history:
                    chat_history = [
                        {"role": "assistant", "content": event.get('stream_text', "")},
                        {"role": "assistant", "content": "Hi! I'm LifeFlow. Tell me if you want to change priorities, add stops, or adjust times."}]
                    current_session.chat_history.append({"role": "assistant", "message": "Hi! I'm LifeFlow...", "time": ""})

                yield (
                    gr.update(visible=False),  # Hide S1
                    gr.update(visible=True),  # Show S2
                    gr.update(visible=False),  # Hide S3
                    "",  # Clear Stream
                    gr.HTML(tasks_html),
                    gr.HTML(summary_html),
                    chat_history,
                    current_session.to_dict()
                )

    def chat_wrapper(self, msg, session_data):
        session = UserSession.from_dict(session_data)
        iterator = self.service.modify_task_chat(msg, session)
        for event in iterator:
            sess = event.get("session", session)
            tasks_html = self.service.generate_task_list_html(sess)
            summary_html = event.get("summary_html", gr.HTML())
            gradio_history = self._get_gradio_chat_history(sess)
            yield (gradio_history, tasks_html, summary_html, sess.to_dict())

    def step3_wrapper(self, session_data):
        session = UserSession.from_dict(session_data)

        # Init variables
        log_content = ""
        report_content = ""

        tasks_html = self.service.generate_task_list_html(session)
        init_dashboard = self._get_dashboard_html('team', 'working', 'Initializing...')

        # HTML Content (No Indentation)
        loading_html = inspect.cleandoc("""
        <div style="text-align: center; padding: 60px 20px; color: #64748b;">
            <div style="font-size: 48px; margin-bottom: 20px; animation: pulse 1.5s infinite;">🧠</div>
            <div style="font-size: 1.2rem; font-weight: 600; color: #334155;">AI Team is analyzing your request...</div>
            <div style="font-size: 0.9rem; margin-top: 8px;">Checking routes, weather, and optimizing schedule.</div>
        </div>
        <style>@keyframes pulse { 0% { transform: scale(1); opacity: 1; } 50% { transform: scale(1.1); opacity: 0.7; } 100% { transform: scale(1); opacity: 1; } }</style>
        """)

        init_log = '<div style="padding: 10px; color: #94a3b8; font-style: italic;">Waiting for agents...</div>'

        yield (init_dashboard, init_log, loading_html, tasks_html, session.to_dict())

        try:
            iterator = self.service.run_step3_team(session)

            for event in iterator:
                sess = event.get("session", session)
                evt_type = event.get("type")

                if evt_type == "error":
                    raise Exception(event.get("message"))

                if evt_type == "report_stream":
                    report_content = event.get("content", "")

                if evt_type == "reasoning_update":
                    agent, status, msg = event.get("agent_status")
                    time_str = datetime.now().strftime('%H:%M:%S')
                    log_entry = f"""
                    <div style="margin-bottom: 8px; border-left: 3px solid #6366f1; padding-left: 10px;">
                        <div style="font-size: 0.75rem; color: #94a3b8;">{time_str} • {agent.upper()}</div>
                        <div style="color: #334155; font-size: 0.9rem;">{msg}</div>
                    </div>
                    """
                    log_content = log_entry + log_content
                    dashboard_html = self._get_dashboard_html(agent, status, msg)
                    log_html = f'<div style="height: 500px; overflow-y: auto; padding: 10px; background: #fff;">{log_content}</div>'
                    current_report = report_content + "\n\n" if report_content else loading_html

                    yield (dashboard_html, log_html, current_report, tasks_html, sess.to_dict())

                if evt_type in ["report_stream", "reasoning_update"]:
                    yield (dashboard_html, log_html, current_report, tasks_html, sess.to_dict())

                if evt_type == "complete":
                    final_report = event.get("report_html", report_content)
                    final_log = f"""
                    <div style="margin-bottom: 8px; border-left: 3px solid #10b981; padding-left: 10px; background: #ecfdf5; padding: 10px;">
                        <div style="font-weight: bold; color: #059669;">✅ Planning Completed</div>
                    </div>
                    {log_content}
                    """
                    final_log_html = f'<div style="height: 500px; overflow-y: auto; padding: 10px; background: #fff;">{final_log}</div>'

                    yield (
                        self._get_dashboard_html('team', 'complete', 'Planning Finished'),
                        final_log_html,
                        final_report,
                        tasks_html,
                        sess.to_dict()
                    )

        except Exception as e:
            error_msg = str(e)
            gr.Warning(f"⚠️ Planning Interrupted: {error_msg}")

            # 這裡我們做一個特殊的 UI 操作：
            # 雖然這個 Wrapper 的 outputs 綁定的是 Step 3 的內部元件
            # 但我們可以透過修改 Step 3 的 "Log" 或 "Report" 區域來顯示「重試按鈕」
            # 或者 (更進階做法) 在 app.py 綁定 outputs 時，就把 Step 2 container 也包進來。

            # 💡 最簡單不改動 outputs 結構的做法：
            # 在 Log 區域顯示紅色錯誤，並提示用戶手動點擊 "Back"

            error_html = f"""
                    <div style="padding: 20px; background: #fee2e2; border: 1px solid #ef4444; border-radius: 8px; color: #b91c1c;">
                        <strong>❌ Error Occurred:</strong><br>{error_msg}<br><br>
                        Please check your API keys or modify tasks, then try again.
                    </div>
                    """

            # 回傳錯誤訊息到 UI，讓使用者知道發生了什麼
            yield (
                self._get_dashboard_html('team', 'error', 'Failed'),  # Dashboard 變紅
                error_html,  # Log 顯示錯誤
                "",
                tasks_html,
                session.to_dict()
            )

    def step4_wrapper(self, session_data):
        session = UserSession.from_dict(session_data)
        result = self.service.run_step4_finalize(session)

        # 🔥 Wrapper 回傳 7 個值
        if result['type'] == 'success':
            return (
                gr.update(visible=False), # 1. Step 3 Hide
                gr.update(visible=True),  # 2. Step 4 Show
                result['summary_tab_html'], # 3. Summary Tab HTML
                result['report_md'],        # 4. Report Tab Markdown
                result['task_list_html'],   # 5. Task List HTML
                result['map_fig'],          # 6. Map Plot
                session.to_dict()           # 7. Session State
            )
        else:
            return (
                gr.update(visible=True), gr.update(visible=False),
                "", "", "", None,
                session.to_dict()
            )

    def save_settings(self, g, w, prov, m_key, m_sel, fast, g_key_in, f_sel, s_data):
        sess = UserSession.from_dict(s_data)

                # 存入 Session
        sess.custom_settings.update({
                    'google_maps_api_key': g,
                    'openweather_api_key': w,
                    'llm_provider': prov,
                    'model_api_key': m_key,     # 主模型 Key
                    'model': m_sel,             # 主模型 ID
                    'enable_fast_mode': fast,   # 🔥 Fast Mode 開關
                    'groq_fast_model': f_sel,
                    'groq_api_key': g_key_in    # 🔥 獨立 Groq Key
        })
        return gr.update(visible=False), sess.to_dict(), "✅ Configuration Saved"
    # ================= Main UI Builder =================

    def build_interface(self):
        container_css = """
                .gradio-container {
                    max-width: 100% !important; 
                    padding: 0;
                    height: 100vh !important;      /* 1. 關鍵：鎖定高度為視窗大小 */
                    overflow-y: auto !important;   /* 2. 關鍵：內容過長時，在內部產生捲軸 */
                }
                """

        with gr.Blocks(title="LifeFlow AI", css=container_css) as demo:
            gr.HTML(get_enhanced_css())

            session_state = gr.State(UserSession().to_dict())

            with gr.Column(elem_classes="step-container"):
                home_btn, settings_btn, doc_btn = create_header() #theme_btn

                stepper = create_progress_stepper(1)
                status_bar = gr.Markdown("Ready", visible=False)

            # STEP 1
            with gr.Group(visible=True, elem_classes="step-container centered-input-container") as step1_container:
                (input_area, s1_stream_output, user_input, auto_loc, loc_group, lat_in, lon_in, analyze_btn) = create_input_form("")

            # STEP 2
            with gr.Group(visible=False, elem_classes="step-container") as step2_container:
                gr.Markdown("### ✅ Review & Refine Tasks")
                with gr.Row(elem_classes="step2-split-view"):
                    with gr.Column(scale=1):
                        with gr.Group(elem_classes="panel-container"):
                            gr.HTML('<div class="panel-header">📋 Your Tasks</div>')
                            with gr.Group(elem_classes="scrollable-content"):
                                task_summary_box = gr.HTML()
                                task_list_box = gr.HTML()
                    with gr.Column(scale=1):
                        with gr.Group(elem_classes="panel-container chat-panel-native"):
                            chatbot = gr.Chatbot(label="AI Assistant", type="messages", height=540, elem_classes="native-chatbot", bubble_full_width=False)
                            with gr.Row(elem_classes="chat-input-row"):
                                chat_input = gr.Textbox(show_label=False, placeholder="Type to modify tasks...", container=False, scale=5, autofocus=True)
                                chat_send = gr.Button("➤", variant="primary", scale=1, min_width=50)
                with gr.Row(elem_classes="action-footer"):
                    back_btn = gr.Button("← Back", variant="secondary", scale=1)
                    plan_btn = gr.Button("🚀 Start Planning", variant="primary", scale=2)

            # STEP 3
            with gr.Group(visible=False, elem_classes="step-container") as step3_container:
                gr.Markdown("### 🤖 AI Team Operations")

                # 1. Agent Dashboard (保持不變)
                with gr.Group(elem_classes="agent-dashboard-container"):
                    agent_dashboard = gr.HTML(value=self._get_dashboard_html())

                # 2. 主要內容區 (左右分欄)
                with gr.Row():
                    # 🔥 左側：主要報告顯示區 (佔 3/4 寬度)
                    with gr.Column(scale=3):
                        with gr.Tabs():
                            with gr.Tab("📝 Full Report"):
                                # 🔥 修正 2: 加入白底容器 (live-report-wrapper)
                                with gr.Group(elem_classes="live-report-wrapper"):
                                    live_report_md = gr.Markdown()

                            with gr.Tab("📋 Task List"):
                                with gr.Group(elem_classes="panel-container"):
                                    with gr.Group(elem_classes="scrollable-content"):
                                        task_list_s3 = gr.HTML()

                    # 🔥 右側：控制區與日誌 (佔 1/4 寬度)
                    with gr.Column(scale=1):
                        # 🔥 修正 1: 將停止按鈕移到這裡 (右側欄位頂部)
                        # variant="stop" 會呈現紅色，更符合 "緊急停止" 的語意
                        cancel_plan_btn = gr.Button("🛑 Stop & Back to Edit", variant="stop")

                        gr.Markdown("### ⚡ Activity Log")
                        with gr.Group(elem_classes="panel-container"):
                            planning_log = gr.HTML(value="Waiting...")

            # STEP 4
            with gr.Group(visible=False, elem_classes="step-container") as step4_container:
                with gr.Row():
                    # Left: Tabs (Summary / Report / Tasks)
                    with gr.Column(scale=1, elem_classes="split-left-panel"):
                        with gr.Tabs():
                            with gr.Tab("📊 Summary"):
                                summary_tab_output = gr.HTML()

                            with gr.Tab("📝 Full Report"):
                                with gr.Group(elem_classes="live-report-wrapper"):
                                    report_tab_output = gr.Markdown()

                            with gr.Tab("📋 Task List"):
                                task_list_tab_output = gr.HTML()

                    # Right: Hero Map
                    with gr.Column(scale=2, elem_classes="split-right-panel"):
                        map_view = gr.HTML(label="Route Map")

            # Modals & Events
            (settings_modal,
             g_key, g_stat,
             w_key, w_stat,
             llm_provider, main_key, main_stat, model_sel,
             fast_mode_chk, groq_model_sel, groq_key, groq_stat,
             close_set, save_set, set_stat) = create_settings_modal()

            # 2. 綁定 Google Maps 測試
            g_key.blur(
                fn=lambda k: "⏳ Checking..." if k else "",  # 先顯示 checking
                inputs=[g_key], outputs=[g_stat]
            ).then(
                fn=APIValidator.test_google_maps,
                inputs=[g_key],
                outputs=[g_stat]
            )

            # 3. OpenWeather 自動驗證
            w_key.blur(
                fn=lambda k: "⏳ Checking..." if k else "",
                inputs=[w_key], outputs=[w_stat]
            ).then(
                fn=APIValidator.test_openweather,
                inputs=[w_key],
                outputs=[w_stat]
            )

            # 4. 綁定 Main Brain 測試
            main_key.blur(
                fn=lambda k: "⏳ Checking..." if k else "",
                inputs=[main_key], outputs=[main_stat]
            ).then(
                fn=APIValidator.test_llm,
                inputs=[llm_provider, main_key, model_sel],
                outputs=[main_stat]
            )

            # 5. Groq 自動驗證
            groq_key.blur(
                fn=lambda k: "⏳ Checking..." if k else "",
                inputs=[groq_key], outputs=[groq_stat]
            ).then(
                fn=lambda k, m: APIValidator.test_llm("Groq", k, m),
                inputs=[groq_key, groq_model_sel],
                outputs=[groq_stat]
            )

            def resolve_groq_visibility(fast_mode, provider):
                """
                綜合判斷 Groq 相關欄位是否該顯示
                回傳: (Model選單狀態, Key輸入框狀態, 狀態訊息狀態)
                """
                # 1. 如果沒開 Fast Mode -> 全部隱藏 + 清空狀態
                if not fast_mode:
                    return (
                        gr.update(visible=False),
                        gr.update(visible=False),
                        gr.update(visible=False)  # 🔥 清空狀態
                    )

                # 2. 如果開了 Fast Mode
                model_vis = gr.update(visible=True)

                # 3. 判斷 Key 是否需要顯示
                is_main_groq = (provider == "Groq")

                if is_main_groq:
                    # 如果主 Provider 是 Groq，隱藏 Key 並清空狀態 (避免誤導)
                    return (
                        model_vis,
                        gr.update(visible=False),
                        gr.update(visible=False)  # 🔥 清空狀態
                    )
                else:
                    # 否則顯示 Key，狀態保持原樣 (不更動)
                    return (
                        model_vis,
                        gr.update(visible=True),
                        gr.update(visible=True)  # 保持原樣
                    )

            fast_mode_chk.change(
                fn=resolve_groq_visibility,
                inputs=[fast_mode_chk, llm_provider],
                outputs=[groq_model_sel, groq_key, groq_stat]  # 🔥 新增 groq_stat
            )

            # --- 事件 B: 當 "Main Provider" 改變時 ---
            def on_provider_change(provider, fast_mode):
                # 1. 更新 Model 下拉選單 (原有邏輯)
                new_choices = MODEL_OPTIONS.get(provider, [])
                default_val = new_choices[0][1] if new_choices else ""
                model_update = gr.Dropdown(choices=new_choices, value=default_val)

                # 2. 清空 Main API Key (原有邏輯)
                key_clear = gr.Textbox(value="")
                stat_clear = gr.Markdown(value="")

                # 3. 重算 Groq 欄位可見性 (包含狀態清空)
                g_model_vis, g_key_vis, g_stat_vis = resolve_groq_visibility(fast_mode, provider)

                # 回傳順序要對應下面的 outputs
                return (
                    model_update,  # model_sel
                    key_clear,  # main_key
                    stat_clear,  # main_stat
                    g_model_vis,  # groq_model_sel
                    g_key_vis,  # groq_key
                    g_stat_vis  # groq_stat (新)
                )

            llm_provider.change(
                fn=on_provider_change,
                inputs=[llm_provider, fast_mode_chk],
                # 🔥 outputs 增加到 6 個，確保所有狀態都被重置
                outputs=[model_sel, main_key, main_stat, groq_model_sel, groq_key, groq_stat]
            )

            doc_modal, close_doc_btn = create_doc_modal()

            settings_btn.click(fn=lambda: gr.update(visible=True), outputs=[settings_modal])
            close_set.click(fn=lambda: gr.update(visible=False), outputs=[settings_modal])
            doc_btn.click(fn=lambda: gr.update(visible=True), outputs=[doc_modal])
            close_doc_btn.click(fn=lambda: gr.update(visible=False), outputs=[doc_modal])
            #theme_btn.click(fn=None, js="() => { document.querySelector('.gradio-container').classList.toggle('theme-dark'); }")

            analyze_btn.click(
                fn=self.analyze_wrapper,
                inputs=[user_input, auto_loc, lat_in, lon_in, session_state,],
                outputs=[step1_container, step2_container, step3_container, s1_stream_output, task_list_box, task_summary_box, chatbot, session_state]
            ).then(fn=lambda: update_stepper(2), outputs=[stepper])

            chat_send.click(fn=self.chat_wrapper, inputs=[chat_input, session_state], outputs=[chatbot, task_list_box, task_summary_box, session_state]).then(fn=lambda: "", outputs=[chat_input])
            chat_input.submit(fn=self.chat_wrapper, inputs=[chat_input, session_state], outputs=[chatbot, task_list_box, task_summary_box, session_state]).then(fn=lambda: "", outputs=[chat_input])

            step3_start = plan_btn.click(
                fn=lambda: (gr.update(visible=False), gr.update(visible=True), update_stepper(3)),
                outputs=[step2_container, step3_container, stepper]
            ).then(
                fn=self.step3_wrapper,
                inputs=[session_state],
                outputs=[agent_dashboard, planning_log, live_report_md, task_list_s3, session_state],
                show_progress="hidden"
            )

            cancel_plan_btn.click(
                fn=self.cancel_wrapper,  # 1. 先執行取消
                inputs=[session_state],
                outputs=None
            ).then(
                fn=lambda: (gr.update(visible=True), gr.update(visible=False), update_stepper(2)),  # 2. 再切換 UI
                inputs=None,
                outputs=[step2_container, step3_container, stepper]
            )

            step3_start.then(
                fn=self.step4_wrapper,
                inputs=[session_state],
                # 🔥🔥🔥 關鍵修正：這裡列出了 7 個 Outputs，必須對應 Wrapper 回傳的 7 個值
                outputs=[
                    step3_container,      # 1. Hide Step 3
                    step4_container,      # 2. Show Step 4
                    summary_tab_output,   # 3. Summary Tab
                    report_tab_output,    # 4. Report Tab
                    task_list_tab_output, # 5. Task List Tab
                    map_view,             # 6. Map
                    session_state         # 7. State
                ]
            ).then(fn=lambda: update_stepper(4), outputs=[stepper])

            def reset_all(session_data):
                old_session = UserSession.from_dict(session_data)
                new_session = UserSession()
                new_session.custom_settings = old_session.custom_settings
                return (gr.update(visible=True), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), update_stepper(1), new_session.to_dict(), "")

            home_btn.click(
                fn=self.cancel_wrapper,  # 1. 先執行取消 (防止後台繼續跑)
                inputs=[session_state],
                outputs=None
            ).then(
                fn=reset_all,  # 2. 再重置 UI
                inputs=[session_state],
                outputs=[step1_container, step2_container, step3_container, step4_container, stepper, session_state,
                         user_input]
            )

            back_btn.click(
                fn=self.cancel_wrapper,
                inputs=[session_state],
                outputs=None
            ).then(
                fn=reset_all,
                inputs=[session_state],
                outputs=[step1_container, step2_container, step3_container, step4_container, stepper, session_state,
                         user_input]
            )

            save_set.click(
                fn=self.save_settings,
                # 輸入參數對應上面的 create_settings_modal 回傳順序
                inputs=[g_key, w_key, llm_provider, main_key, model_sel, fast_mode_chk, groq_key, groq_model_sel, session_state],
                outputs=[settings_modal, session_state, status_bar]
            )



            auto_loc.change(fn=toggle_location_inputs, inputs=auto_loc, outputs=loc_group)
            demo.load(fn=self._check_api_status, inputs=[session_state], outputs=[status_bar])

        return demo

def main():
    app = LifeFlowAI()
    demo = app.build_interface()
    demo.launch(server_name="0.0.0.0", server_port=8080, share=True, show_error=True)
    #7860
if __name__ == "__main__":
    main()