import gradio as gr
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
from i18n import get_text

# --- 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, lang):
    """Calculate word count and read time."""
    if not text:
        return get_text("writer_stats_default", lang)
    words = len(text.split())
    read_time = max(1, words // 200) # Average reading speed
    return get_text("writer_stats_format", lang).format(
        words=words, read_time=read_time
    )

# --- UI Construction ---

def create_smart_writer_tab(lang_state: gr.State):
    lang = lang_state.value
    debounce_manager = DebounceManager(debounce_time=2.0, tick_time=1.0, loading_text=get_text("writer_debounce_loading_text", lang))

    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=get_text("writer_style_input_label", lang),
                lines=8, 
                value=MOCK_STYLE,
                interactive=True
            )
            
            with gr.Accordion(
                get_text("writer_kb_accordion_label", lang), open=True
            ) as kb_accordion:
                kb_input = gr.Dataframe(
                    headers=[get_text("writer_kb_dataframe_headers", lang), '?'],
                    datatype=["str", "str"],
                    value=MOCK_KNOWLEDGE_BASE,
                    interactive=True,
                    wrap=True
                )
                with gr.Row():
                    btn_suggest_kb = gr.Button(get_text("writer_suggest_kb_button", lang), size="sm")
                
                suggested_kb_dataframe = gr.Dataframe(
                    headers=["Term", "Description"],
                    datatype=["str", "str"],
                    visible=False, 
                    interactive=False,
                    label=get_text("writer_suggested_kb_label", lang),
                )

            with gr.Accordion(get_text("writer_short_outline_accordion_label", lang), open=True) as short_outline_accordion:
                
                short_outline_input = gr.Dataframe(
                    headers=[get_text("writer_short_outline_dataframe_headers", lang), '?'],
                    datatype=["bool", "str"],
                    value=MOCK_SHORT_TERM_OUTLINE,
                    interactive=True,
                    label="???",
                    col_count=(2, "fixed"),
                )
                with gr.Row():
                    btn_sync_outline = gr.Button(get_text("writer_sync_outline_button", lang), size="sm")

            with gr.Accordion(get_text("writer_long_outline_accordion_label", lang), open=False) as long_outline_accordion:
                long_outline_input = gr.Dataframe(
                    headers=[get_text("writer_short_outline_dataframe_headers", lang), '?'],
                    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=get_text("writer_flow_suggestion_label", lang),
                        placeholder=get_text("writer_flow_suggestion_placeholder", lang),
                        lines=3,
                        interactive=False,
                        elem_classes=["flow-suggestion-box"],
                    )

                    btn_accept_flow = gr.Button(get_text("writer_accept_flow_button", lang),size="sm",variant="primary",elem_id="btn-action-accept-flow")
                    btn_change_flow = gr.Button(get_text("writer_change_flow_button", lang),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=get_text("writer_inspiration_prompt_label", lang),
                        placeholder=get_text("writer_inspiration_prompt_placeholder", lang),
                        lines=2
                    )
                    prompt_suggestions_dataset = gr.Dataset(
                        label=get_text("writer_prompt_suggestions_label", lang),
                        components=[gr.Textbox(visible=False)],
                        samples=[["生成建议..."], ["生成建议..."], ["生成建议..."]],
                        type="values"
                    )
                    refresh_suggestions_btn = gr.Button(
                        get_text("writer_refresh_suggestions_button", lang),
                        size="sm",
                        variant="secondary",
                    )
                    with gr.Row():
                        btn_generate_para = gr.Button(get_text("writer_generate_para_button", lang),size="sm",variant="primary",elem_id="btn-action-create-paragraph")
                        btn_change_para = gr.Button(get_text("writer_change_para_button", lang),size="sm")
                        btn_accept_para = gr.Button(get_text("writer_accept_para_button", lang),size="sm")

                    para_suggestion_display = gr.Textbox(
                        show_label=False,
                        placeholder=get_text("writer_para_suggestion_placeholder", lang),
                        lines=3,
                        interactive=False
                    )
                
                # Area 3: Adjust/Polish (Placeholder)
                with gr.Column(scale=1, min_width=200):
                    polish_title = gr.Markdown(get_text("writer_polish_title", lang))
                    polish_soon = gr.Markdown(get_text("writer_coming_soon", lang))

            # --- TOOLBAR ---
            with gr.Row(elem_classes=["toolbar"]):
                # --- EDITOR ---
                stats_display = gr.Markdown(get_text("writer_stats_default", lang))

            editor = gr.Textbox(
                label=get_text("writer_editor_label", lang),
                placeholder=get_text("writer_editor_placeholder", lang),
                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, lang_state], 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]
    )

    return {
        "style_input": style_input,
        "kb_accordion": kb_accordion,
        "kb_input": kb_input,
        "btn_suggest_kb": btn_suggest_kb,
        "suggested_kb_dataframe": suggested_kb_dataframe,
        "short_outline_accordion": short_outline_accordion,
        "short_outline_input": short_outline_input,
        "btn_sync_outline": btn_sync_outline,
        "long_outline_accordion": long_outline_accordion,
        "long_outline_input": long_outline_input,
        "flow_suggestion_display": flow_suggestion_display,
        "btn_accept_flow": btn_accept_flow,
        "btn_change_flow": btn_change_flow,
        "inspiration_prompt_input": inspiration_prompt_input,
        "prompt_suggestions_dataset": prompt_suggestions_dataset,
        "refresh_suggestions_btn": refresh_suggestions_btn,
        "btn_generate_para": btn_generate_para,
        "btn_change_para": btn_change_para,
        "btn_accept_para": btn_accept_para,
        "para_suggestion_display": para_suggestion_display,
        "polish_title": polish_title,
        "polish_soon": polish_soon,
        "stats_display": stats_display,
        "editor": editor,
    }


def update_language(lang: str, components: dict):
    return {
        components["style_input"]: gr.update(
            label=get_text("writer_style_input_label", lang)
        ),
        components["kb_accordion"]: gr.update(
            label=get_text("writer_kb_accordion_label", lang)
        ),
        components["kb_input"]: gr.update(
            headers=get_text("writer_kb_dataframe_headers", lang)
        ),
        components["btn_suggest_kb"]: gr.update(
            value=get_text("writer_suggest_kb_button", lang)
        ),
        components["suggested_kb_dataframe"]: gr.update(
            label=get_text("writer_suggested_kb_label", lang)
        ),
        components["short_outline_accordion"]: gr.update(
            label=get_text("writer_short_outline_accordion_label", lang)
        ),
        components["short_outline_input"]: gr.update(
            headers=get_text("writer_short_outline_dataframe_headers", lang)
        ),
        components["btn_sync_outline"]: gr.update(
            value=get_text("writer_sync_outline_button", lang)
        ),
        components["long_outline_accordion"]: gr.update(
            label=get_text("writer_long_outline_accordion_label", lang)
        ),
        components["long_outline_input"]: gr.update(
            headers=get_text("writer_short_outline_dataframe_headers", lang)
        ),  # Assuming same headers
        components["flow_suggestion_display"]: gr.update(
            label=get_text("writer_flow_suggestion_label", lang),
            placeholder=get_text("writer_flow_suggestion_placeholder", lang),
        ),
        components["btn_accept_flow"]: gr.update(
            value=get_text("writer_accept_flow_button", lang)
        ),
        components["btn_change_flow"]: gr.update(
            value=get_text("writer_change_flow_button", lang)
        ),
        components["inspiration_prompt_input"]: gr.update(
            label=get_text("writer_inspiration_prompt_label", lang),
            placeholder=get_text("writer_inspiration_prompt_placeholder", lang),
        ),
        components["prompt_suggestions_dataset"]: gr.update(
            label=get_text("writer_prompt_suggestions_label", lang)
        ),
        components["refresh_suggestions_btn"]: gr.update(
            value=get_text("writer_refresh_suggestions_button", lang)
        ),
        components["btn_generate_para"]: gr.update(
            value=get_text("writer_generate_para_button", lang)
        ),
        components["btn_change_para"]: gr.update(
            value=get_text("writer_change_para_button", lang)
        ),
        components["btn_accept_para"]: gr.update(
            value=get_text("writer_accept_para_button", lang)
        ),
        components["para_suggestion_display"]: gr.update(
            placeholder=get_text("writer_para_suggestion_placeholder", lang)
        ),
        components["polish_title"]: gr.update(
            value=get_text("writer_polish_title", lang)
        ),
        components["polish_soon"]: gr.update(
            value=get_text("writer_coming_soon", lang)
        ),
        components["editor"]: gr.update(
            label=get_text("writer_editor_label", lang),
            placeholder=get_text("writer_editor_placeholder", lang),
        ),
        components["stats_display"]: gr.update(
            value=get_text("writer_stats_default", lang)
        ),  # Reset stats on lang change
    }