tab_smart_writer.py

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
from smart_writer_kit.agent_for_paragraph_continuation import fetch_paragraph_continuation_agent
from smart_writer_kit.agent_for_prompt_suggestion import fetch_prompt_suggestions_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
from ui_components.debounce_manager import DebounceManager

# --- 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():
    # Initialize DebounceManager
    debounce_manager = DebounceManager(debounce_time=2.0, tick_time=0.3, loading_text="稍后开始续写")

    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):

                    flow_suggestion_display = gr.Textbox(
                        show_label=True,
                        label="实时续写建议",
                        placeholder="（等待输入或点击“换一个”...）",
                        lines=3,
                        interactive=False,
                        elem_classes=["flow-suggestion-box"],
                    )

                    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')

                    # Debounce Progress Indicator (Using Manager)
                    debounce_state, debounce_timer, debounce_progress = debounce_manager.create_ui()
                    debounce_progress.visible = True 

                # Area 2: Paragraph Continuation (Inspiration)
                with gr.Column(scale=1, min_width=200):
                    inspiration_prompt_input = gr.Textbox(
                        label="续写提示",
                        placeholder="例如：写一段关于...的描写",
                        lines=2
                    )
                    
                    prompt_suggestions_dataset = gr.Dataset(
                        label="推荐提示 (点击填入)",
                        components=[gr.Textbox(visible=False)],
                        samples=[["生成建议..."], ["生成建议..."], ["生成建议..."]],
                        type="values"
                    )
                    
                    refresh_suggestions_btn = gr.Button("🎲 换一批建议", size="sm", variant="secondary") # Combined trigger

                    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 (Using DebounceManager)
    
    # Bind reset logic to editor change
    editor.change(
        fn=debounce_manager.reset, 
        inputs=[editor, style_input, kb_input, short_outline_input, long_outline_input], # Capture all context as payload
        outputs=[debounce_state, debounce_timer, debounce_progress]
    )

    # Bind tick logic
    def flow_suggestion_trigger(editor_content, style, kb, short_outline, long_outline):
         return fetch_flow_suggestion_agent(editor_content, style, kb, short_outline, long_outline)

    # Note: debounce_manager.tick calls the trigger function. 
    # The lambda is used to pass the specific trigger function for this tab.
    debounce_timer.tick(
        fn=lambda s: debounce_manager.tick(s, flow_suggestion_trigger), 
        inputs=[debounce_state], 
        outputs=[debounce_progress, debounce_state, debounce_timer, flow_suggestion_display]
    )
    
    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]
    )

    # 3. Paragraph Continuation Logic (Updated with prompt input)
    def generate_paragraph_wrapper(prompt_val, editor_val, style, kb, short, long_):
        return fetch_paragraph_continuation_agent(prompt_val, editor_val, style, kb, short, long_)

    for btn in [btn_generate_para, btn_change_para]:
        btn.click(
            fn=generate_paragraph_wrapper,
            inputs=[inspiration_prompt_input, 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]
    )
    
    # Suggestions Logic
    # Trigger for suggestion generation
    def refresh_suggestions_wrapper(editor_content, style, kb, short_outline, long_outline):
        s1, s2, s3 = fetch_prompt_suggestions_agent(editor_content, style, kb, short_outline, long_outline)
        # Return a gr.update object to properly update the Dataset component
        return gr.update(samples=[[s1], [s2], [s3]])

    refresh_suggestions_btn.click(
        fn=refresh_suggestions_wrapper,
        inputs=[editor, style_input, kb_input, short_outline_input, long_outline_input],
        outputs=[prompt_suggestions_dataset]
    )
    
    # Dataset click -> fill prompt input
    def fill_prompt_from_dataset(val):
        return val[0]

    prompt_suggestions_dataset.click(
        fn=fill_prompt_from_dataset,
        inputs=prompt_suggestions_dataset,
        outputs=inspiration_prompt_input
    )

    # 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]
    )