import gradio as gr
import time
from smart_writer_kit.agent_for_streaming_completion import fetch_flow_suggestion_agent, accept_flow_suggestion_agent
from smart_writer_kit.agent_for_inspiration_expansion import fetch_inspiration_agent, apply_inspiration_agent, fetch_paragraph_continuation_agent
from smart_writer_kit.agent_for_outline_update import update_outline_status_agent
from smart_writer_kit.agent_for_kb_update import suggest_new_kb_terms_agent

# --- Mock Data (for UI population only) ---

MOCK_STYLE = """故事：人类逐渐走向消亡时，人形机器人的休闲生活。
风格：自然平淡，文字细腻，描绘未来的荒凉与宁静交织的景象。
主题：探索人类与机器的界限，记忆与身份的意义。
"""

MOCK_KNOWLEDGE_BASE = [
    ["Alpha", "故事的主角，女性人形机器人，外表与人类无异。性格有线"],
    ["横滨", "故事发生的主要城市。由于海平面上升，城市部分地区被淹没，形成独特的水上景观。"]
]

MOCK_SHORT_TERM_OUTLINE = [
    [False, "故事的场景设定：海平面上升后的城市景观。"],
    [False, "介绍主角 Alpha 的日常生活和她与其他机器人的互动。"],
    [False, "Alpha 发现了一张旧照片，勾起了她对过去人类生活的好奇心。"],
    [False, "奶油蛋糕的制作方法。"]
]

## 按日常向动画剧情走向写的长纲要。具体。
MOCK_LONG_TERM_OUTLINE = [
    [False, "介绍故事背景。人类逐渐减少，机器人和人的互动。"],
    [False, "Alpha 决定离开居住地，到东京寻找失散的朋友。"],
    [False, "月亮变成了一个巨大的 Disco 灯球。机器人不受控制地开始跳舞，导致全球范围内的混乱。"],
]

# --- UI Helper Functions ---
def get_stats(text):
    """Calculate word count and read time."""
    if not text:
        return "0 Words | 0 mins"
    words = len(text.split())
    read_time = max(1, words // 200) # Average reading speed
    return f"{words} Words | ~{read_time} mins"

# --- UI Construction ---

def create_smart_writer_tab():
    debounce_state = gr.State({"last_change": 0, "active": False, "style": "", "kb": [], "short_outline": [], "long_outline": []})
    debounce_timer = gr.Timer(0.1, active=True)
    

    with gr.Row(equal_height=False, elem_id="indicator-writing-tab"):
        # --- Left Column: Entity Console ---
        with gr.Column(scale=1) as left_panel:
            
            style_input = gr.Textbox(
                label="整体故事和风格", 
                lines=8, 
                value=MOCK_STYLE,
                interactive=True
            )
            
            with gr.Accordion("写作知识库", open=True):
                
                kb_input = gr.Dataframe(
                    headers=["名称", "说明"],
                    datatype=["str", "str"],
                    value=MOCK_KNOWLEDGE_BASE,
                    interactive=True,
                    wrap=True
                )
                with gr.Row():
                    btn_suggest_kb = gr.Button("🔍 提取新词条", size="sm")
                
                suggested_kb_dataframe = gr.Dataframe(
                    headers=["Term", "Description"],
                    datatype=["str", "str"],
                    visible=False, 
                    interactive=False,
                    label="推荐词条"
                )

            with gr.Accordion("当前章节大纲", open=True):
                
                short_outline_input = gr.Dataframe(
                    headers=["Done", "Task"],
                    datatype=["bool", "str"],
                    value=MOCK_SHORT_TERM_OUTLINE,
                    interactive=True,
                    label="当前章节大纲",
                    col_count=(2, "fixed"),
                )
                with gr.Row():
                    btn_sync_outline = gr.Button("🔄 同步状态", size="sm")

            with gr.Accordion("故事整体大纲", open=False):
                long_outline_input = gr.Dataframe(
                    headers=["Done", "Task"],
                    datatype=["bool", "str"],
                    value=MOCK_LONG_TERM_OUTLINE,
                    interactive=True,
                    label="故事总纲",
                    col_count=(2, "fixed"),
                )

        # --- Right Column: Writing Canvas ---
        with gr.Column(scale=5):
            
            # --- RIBBON AREA (Top of Editor) ---
            with gr.Row(variant="panel", elem_classes=["ribbon-container"]):
                
                # Area 1: Real-time Continuation (Flow)
                with gr.Column(scale=1, min_width=200):
                    gr.Markdown("#### ⚡️ 实时续写")
                    with gr.Row():
                        btn_accept_flow = gr.Button("采纳续写 (Tab)", size="sm", variant="primary", elem_id='btn-action-accept-flow')
                        btn_change_flow = gr.Button("换一个 (Shift+Tab)", size="sm", elem_id='btn-action-change-flow')
                    
                    flow_suggestion_display = gr.Textbox(
                        show_label=False,
                        placeholder="（等待输入或点击“换一个”...）",
                        lines=3,
                        interactive=False,
                        elem_classes=["flow-suggestion-box"]
                    )
                    # Debounce Progress Indicator
                    debounce_progress = gr.HTML(value="", visible=False)

                # Area 2: Paragraph Continuation (Inspiration)
                with gr.Column(scale=1, min_width=200):
                    gr.Markdown("#### ✨ 整段续写")
                    with gr.Row():
                        btn_generate_para = gr.Button("整段续写 (Cmd+Enter)", size="sm", variant="primary", elem_id="btn-action-create-paragraph")
                        btn_change_para = gr.Button("换一个", size="sm")
                        btn_accept_para = gr.Button("采纳", size="sm")

                    para_suggestion_display = gr.Textbox(
                        show_label=False,
                        placeholder="（点击“整段续写”生成内容...）",
                        lines=3,
                        interactive=False
                    )
                
                # Area 3: Adjust/Polish (Placeholder)
                with gr.Column(scale=1, min_width=200):
                    gr.Markdown("#### 🛠️ 调整润色")
                    gr.Markdown("(Coming Soon)")

            # --- TOOLBAR ---
            with gr.Row(elem_classes=["toolbar"]):
                stats_display = gr.Markdown("0 Words | 0 mins")
            
            # --- EDITOR ---
            editor = gr.Textbox(
                label="沉浸写作画布", 
                placeholder="开始你的创作...", 
                lines=25, # Reduced lines slightly to accommodate ribbon
                elem_classes=["writing-editor"],
                elem_id="writing-editor",
                show_label=False,
            )


    # --- Interactions ---

    # 1. Stats
    editor.change(fn=get_stats, inputs=editor, outputs=stats_display)

    # 2. Flow Suggestion Logic
    def start_debounce(editor_content, style, kb, short_outline, long_outline):
        return {"last_change": time.time(), "active": True, "style": style, "kb": kb, "short_outline": short_outline, "long_outline": long_outline}, gr.update(active=True), gr.update(visible=True, value="<progress value='0' max='100'></progress> 补全中... 3.0s")

    def update_debounce(debounce_state, editor_content):
        if not debounce_state["active"]:
            return gr.update(), gr.update(), debounce_state, gr.update()
        elapsed = time.time() - debounce_state["last_change"]
        if elapsed >= 3:
            suggestion = fetch_flow_suggestion_agent(editor_content, debounce_state["style"], debounce_state["kb"], debounce_state["short_outline"], debounce_state["long_outline"])
            return gr.update(visible=False), suggestion, {"last_change": 0, "active": False, "style": "", "kb": [], "short_outline": [], "long_outline": []}, gr.update(active=False)
        else:
            progress = int((elapsed / 3) * 100)
            remaining = 3 - elapsed
            progress_html = f"<progress value='{progress}' max='100'></progress> 补全中... {remaining:.1f}s"
            return gr.update(value=progress_html), gr.update(), debounce_state, gr.update()

    editor.change(fn=start_debounce, inputs=[editor, style_input, kb_input, short_outline_input, long_outline_input], outputs=[debounce_state, debounce_timer, debounce_progress])
    debounce_timer.tick(fn=update_debounce, inputs=[debounce_state, editor], outputs=[debounce_progress, flow_suggestion_display, debounce_state, debounce_timer])
    
    btn_change_flow.click(fn=fetch_flow_suggestion_agent, inputs=[editor, style_input, kb_input, short_outline_input, long_outline_input], outputs=flow_suggestion_display)

    accept_flow_fn_inputs = [editor, flow_suggestion_display]
    # accept_flow_suggestion_agent returns modified editor text
    btn_accept_flow.click(
        fn=lambda e, s: (accept_flow_suggestion_agent(e, s), ""), # Accept and clear suggestion
        inputs=accept_flow_fn_inputs, 
        outputs=[editor, flow_suggestion_display], 
        show_progress="hidden"
    )

    # 3. Paragraph Continuation Logic
    def generate_paragraph_wrapper(editor_val, style, kb, short, long_):
        return fetch_paragraph_continuation_agent(None, editor_val, style, kb, short, long_)

    for btn in [btn_generate_para, btn_change_para]:
        btn.click(
            fn=generate_paragraph_wrapper,
            inputs=[editor, style_input, kb_input, short_outline_input, long_outline_input],
            outputs=[para_suggestion_display]
        )
    
    def accept_para_wrapper(curr, new):
        # Reuse apply_inspiration_agent but extract text. 
        # It returns (new_text, modal_update, empty_string)
        res = apply_inspiration_agent(curr, new)
        return res[0], ""

    btn_accept_para.click(
        fn=accept_para_wrapper,
        inputs=[editor, para_suggestion_display],
        outputs=[editor, para_suggestion_display]
    )

    # 4. Agent-based Context Updates
    btn_sync_outline.click(
        fn=update_outline_status_agent,
        inputs=[short_outline_input, editor],
        outputs=[short_outline_input]
    )
    btn_suggest_kb.click(
        fn=suggest_new_kb_terms_agent,
        inputs=[kb_input, editor],
        outputs=[suggested_kb_dataframe]
    )