app.py

from __future__ import annotations

from typing import Dict, Iterator, List, Tuple

import gradio as gr

from chat_handlers import ask_report_chat, clear_report_chat
from config import CATALOG_PATH, CUSTOM_CSS, QUESTIONS_PATH, TASK_GROUPS
from helpers import load_question_bank, load_state_catalog
from llm_handlers import generate_play_plan_output, generate_task_output
from models import StateInfo
from vector_store import load_vector_store



def update_task_type_options(task_group: str):
    """Update task type options based on selected task group."""
    options = TASK_GROUPS.get(task_group, [])
    value = options[0] if options else None
    return gr.Dropdown(choices=options, value=value)


def update_play_options(state_label: str):
    """Update playbook options based on selected state."""
    if not state_label or state_label not in STATES_BY_LABEL:
        return gr.Dropdown(choices=[], value=None)
    state = STATES_BY_LABEL[state_label]
    plays = STATE_PLAYBOOKS.get(state.slug) or state.plays
    value = plays[0] if plays else None
    return gr.Dropdown(choices=plays, value=value)


def toggle_online_focus(is_enabled: bool):
    """Toggle visibility of online focus textbox."""
    return gr.update(visible=is_enabled)


def workflow_markdown() -> str:
    """Generate workflow documentation markdown."""
    state_list = ", ".join(state.label for state in STATES) if STATES else "No indexed states yet"
    return f"""
### Workflow
1. Use **Report Q&A Chat** to ask plain-language questions about selected state reports.
2. Use **Marketing Task Studio** for structured task generation (multi-state allowed).
3. Use **Playbook Planner** for a selected playbook (single state).


### Current indexed states
- {state_list}

### Notes
- Be patient it takes 1 to 3 minutes to generate reports as the model retrieves and reasons over multiple report snippets.
- All outputs are grounded in report retrieval.
- `Search online` adds online context from trusted or official domains where available.
"""


# Wrapper functions to inject dependencies into handlers
def _ask_report_chat_wrapper(
    message: str,
    ui_history: List[dict],
    state_labels: List[str],
    search_online: bool,
    online_focus: str,
    request: gr.Request,
):
    """Wrapper to inject dependencies into ask_report_chat."""
    return ask_report_chat(
        message=message,
        ui_history=ui_history,
        state_labels=state_labels,
        search_online=search_online,
        online_focus=online_focus,
        request=request,
        vector_store=VECTOR_STORE,
        vector_error=VECTOR_ERROR,
        states_by_label=STATES_BY_LABEL,
        states_by_slug=STATES_BY_SLUG,
        question_bank=QUESTION_BANK,
    )


def _generate_task_output_wrapper(
    task_group: str,
    task_type: str,
    state_labels: List[str],
    objective: str,
    audience_notes: str,
    channels: List[str],
    tone: str,
    output_format: str,
    idea_count: int,
    include_kpis: bool,
    search_online: bool,
    online_focus: str,
    temperature: float,
):
    """Wrapper to inject dependencies into generate_task_output."""
    return generate_task_output(
        task_group=task_group,
        task_type=task_type,
        state_labels=state_labels,
        objective=objective,
        audience_notes=audience_notes,
        channels=channels,
        tone=tone,
        output_format=output_format,
        idea_count=idea_count,
        include_kpis=include_kpis,
        search_online=search_online,
        online_focus=online_focus,
        temperature=temperature,
        vector_store=VECTOR_STORE,
        vector_error=VECTOR_ERROR,
        states_by_label=STATES_BY_LABEL,
        states_by_slug=STATES_BY_SLUG,
    )


def _generate_play_plan_output_wrapper(
    state_label: str,
    play_title: str,
    objective: str,
    constraints: str,
    horizon: str,
    budget_level: str,
    search_online: bool,
    online_focus: str,
    temperature: float,
):
    """Wrapper to inject dependencies into generate_play_plan_output."""
    return generate_play_plan_output(
        state_label=state_label,
        play_title=play_title,
        objective=objective,
        constraints=constraints,
        horizon=horizon,
        budget_level=budget_level,
        search_online=search_online,
        online_focus=online_focus,
        temperature=temperature,
        vector_store=VECTOR_STORE,
        vector_error=VECTOR_ERROR,
        states_by_label=STATES_BY_LABEL,
    )


def _generate_task_output_stream(
    task_group: str,
    task_type: str,
    state_labels: List[str],
    objective: str,
    audience_notes: str,
    channels: List[str],
    tone: str,
    output_format: str,
    idea_count: int,
    include_kpis: bool,
    search_online: bool,
    online_focus: str,
    temperature: float,
) -> Iterator[Tuple[str, str, str, dict]]:
    """Stream task generation progress to UI before final content arrives."""
    yield (
        "",
        "Generating report... this usually takes 1 to 3 minutes.",
        "Collecting report snippets and source trace...",
        gr.update(interactive=False),
    )
    try:
        report_code, report_markdown, source_markdown = _generate_task_output_wrapper(
            task_group=task_group,
            task_type=task_type,
            state_labels=state_labels,
            objective=objective,
            audience_notes=audience_notes,
            channels=channels,
            tone=tone,
            output_format=output_format,
            idea_count=idea_count,
            include_kpis=include_kpis,
            search_online=search_online,
            online_focus=online_focus,
            temperature=temperature,
        )
    except Exception as exc:
        report_code = ""
        report_markdown = f"Generation failed: {exc}"
        source_markdown = ""
    yield (
        report_code,
        report_markdown,
        source_markdown,
        gr.update(interactive=True),
    )


def _generate_play_plan_output_stream(
    state_label: str,
    play_title: str,
    objective: str,
    constraints: str,
    horizon: str,
    budget_level: str,
    search_online: bool,
    online_focus: str,
    temperature: float,
) -> Iterator[Tuple[str, str, str, dict]]:
    """Stream play-plan generation progress to UI before final content arrives."""
    yield (
        "",
        "Generating play plan... this usually takes 1 to 3 minutes.",
        "Collecting report snippets and source trace...",
        gr.update(interactive=False),
    )
    try:
        report_code, report_markdown, source_markdown = _generate_play_plan_output_wrapper(
            state_label=state_label,
            play_title=play_title,
            objective=objective,
            constraints=constraints,
            horizon=horizon,
            budget_level=budget_level,
            search_online=search_online,
            online_focus=online_focus,
            temperature=temperature,
        )
    except Exception as exc:
        report_code = ""
        report_markdown = f"Generation failed: {exc}"
        source_markdown = ""
    yield (
        report_code,
        report_markdown,
        source_markdown,
        gr.update(interactive=True),
    )


def build_app() -> gr.Blocks:
    default_group = next(iter(TASK_GROUPS))
    default_task = TASK_GROUPS[default_group][0]
    state_labels = [state.label for state in STATES]
    default_state_multi = [state_labels[0]] if state_labels else []
    default_state_single = state_labels[0] if state_labels else None

    default_plays = []
    if default_state_single and default_state_single in STATES_BY_LABEL:
        state = STATES_BY_LABEL[default_state_single]
        default_plays = STATE_PLAYBOOKS.get(state.slug) or state.plays

    with gr.Blocks(title="FORTIFIED Marketing Copilot") as demo:
        gr.HTML(
            """
            <div id="hero">
              <h1>FORTIFIED Marketing Copilot</h1>
              <p>Behavioral marketing generation for state-level FORTIFIED roof adoption plans.</p>
            </div>
            """
        )

        if VECTOR_ERROR:
            gr.Markdown(f"⚠️ **Knowledge base warning:** {VECTOR_ERROR}")
        if QUESTION_BANK_ERROR:
            gr.Markdown(f"⚠️ **Survey chart warning:** {QUESTION_BANK_ERROR}")

        with gr.Tab("Instructions"):
            gr.Markdown(workflow_markdown())

        with gr.Tab("Report Q&A Chat"):
            with gr.Row(equal_height=False):
                with gr.Column(scale=4, elem_classes=["card"]):
                    chat_states = gr.Dropdown(
                        choices=state_labels,
                        value=default_state_multi,
                        multiselect=True,
                        label="State reports (select one or more)",
                    )
                    chat_search_online = gr.Checkbox(value=False, label="Search online (selected websites)")
                    chat_online_focus = gr.Textbox(
                        label="Online search focus (optional)",
                        placeholder="e.g., latest homeowner grant updates and insurer discount guidance",
                        visible=False,
                        value="",
                    )
                    clear_chat = gr.Button("Clear chat")
                    related_charts = gr.Gallery(
                        value=[],
                        label="Related survey charts",
                        elem_id="related-charts-gallery",
                        columns=1,
                        height=460,
                        object_fit="contain",
                        show_label=True,
                        preview=True,
                    )
                    chart_debug = gr.Markdown("`Chart debug (latest response)` qids=[] | images=[] | shown=0",
                                              visible=False)

                with gr.Column(scale=5, elem_classes=["card"]):
                    chatbox = gr.Chatbot(
                        value=[],
                        # type="messages",
                        label="Report assistant",
                        height=420,
                    )
                    chat_input = gr.Textbox(
                        label=f"Ask about the selected reports (shift + ↵  to send)",
                        placeholder="e.g., What are the top barriers for homeowner adoption and what messages should we test first?",
                        lines=2,
                    )
                    ask_chat = gr.Button("Ask", variant="primary")
                    # chat_sources = gr.Code(label="Source trace", language="markdown", lines=10, show_label=True)
                    chat_sources = gr.Markdown(label="Source trace", value="Sources will appear here.", padding=True)
            chat_search_online.change(fn=toggle_online_focus, inputs=chat_search_online, outputs=chat_online_focus, show_progress="hidden")
            ask_chat.click(
                fn=_ask_report_chat_wrapper,
                inputs=[chat_input, chatbox, chat_states, chat_search_online, chat_online_focus],
                outputs=[chat_input, chatbox, chat_sources, related_charts, chart_debug],
            )
            chat_input.submit(
                fn=_ask_report_chat_wrapper,
                inputs=[chat_input, chatbox, chat_states, chat_search_online, chat_online_focus],
                outputs=[chat_input, chatbox, chat_sources, related_charts, chart_debug],
            )
            clear_chat.click(
                fn=clear_report_chat,
                outputs=[chat_input, chatbox, chat_sources, related_charts, chart_debug],
            )

        with gr.Tab("Marketing Task Studio"):
            with gr.Row(equal_height=False):
                with gr.Column(scale=4, elem_classes=["card"]):
                    task_group = gr.Radio(
                        choices=list(TASK_GROUPS.keys()),
                        value=default_group,
                        label="Task family",
                    )
                    task_type = gr.Dropdown(
                        choices=TASK_GROUPS[default_group],
                        value=default_task,
                        label="Task type",
                    )
                    states_multi = gr.Dropdown(
                        choices=state_labels,
                        value=default_state_multi,
                        multiselect=True,
                        label="State reports (select one or more)",
                    )
                    objective = gr.Textbox(
                        label="Business objective",
                        placeholder="e.g., Increase qualified FORTIFIED leads by 25% before storm season",
                        lines=2,
                    )
                    audience = gr.Textbox(
                        label="Audience details",
                        placeholder="e.g., Homeowners with roofs 11+ years old and high paperwork friction",
                        lines=2,
                    )
                    channels = gr.CheckboxGroup(
                        choices=[
                            "Paid Search",
                            "SEO",
                            "Email",
                            "Facebook",
                            "Instagram",
                            "YouTube",
                            "Direct Mail",
                            "Contractor Co-Marketing",
                        ],
                        label="Priority channels",
                    )
                    with gr.Row():
                        tone = gr.Dropdown(
                            choices=[
                                "Authoritative",
                                "Local and practical",
                                "Data-driven",
                                "Reassuring",
                                "Urgent but calm",
                            ],
                            value="Local and practical",
                            label="Brand tone",
                        )
                        output_format = gr.Dropdown(
                            choices=["Structured brief", "Campaign table", "90-day roadmap"],
                            value="Structured brief",
                            label="Output format",
                        )
                    with gr.Row():
                        idea_count = gr.Slider(minimum=3, maximum=12, value=6, step=1, label="Idea count")
                        temperature = gr.Slider(minimum=0.0, maximum=1.0, value=0.6, step=0.05, label="Creativity")
                    include_kpis = gr.Checkbox(value=True, label="Include KPI framework")
                    search_online = gr.Checkbox(value=False, label="Search online (selected websites)")
                    online_focus = gr.Textbox(
                        label="Online search focus (optional)",
                        placeholder="e.g., latest Louisiana grant updates and insurer discount guidance",
                        visible=False,
                        value="",
                    )
                    generate_task = gr.Button("Generate", variant="primary")

                with gr.Column(scale=5, ):
                    task_output_m = gr.Markdown(label="Report", value="Generated output will appear here.", padding=True)
                    task_output = gr.Code(label="Report", value="Generated output will appear here.",language="markdown", lines=13, show_label=True, visible=False,)
                    gr.Markdown("---")
                    task_sources = gr.Markdown(label="Source trace", value="Sources will appear here.", padding=True)

            task_group.change(update_task_type_options, inputs=task_group, outputs=task_type)
            search_online.change(fn=toggle_online_focus, inputs=search_online, outputs=online_focus, show_progress="hidden")
            generate_task.click(
                fn=_generate_task_output_stream,
                inputs=[
                    task_group,
                    task_type,
                    states_multi,
                    objective,
                    audience,
                    channels,
                    tone,
                    output_format,
                    idea_count,
                    include_kpis,
                    search_online,
                    online_focus,
                    temperature,
                ],
                outputs=[task_output, task_output_m, task_sources, generate_task],
                show_progress="full",
            )

        with gr.Tab("Playbook Planner"):
            with gr.Row(equal_height=False):
                with gr.Column(scale=4, elem_classes=["card"]):
                    state_single = gr.Dropdown(
                        choices=state_labels,
                        value=default_state_single,
                        multiselect=False,
                        label="State (single selection required)",
                    )
                    play_selector = gr.Dropdown(
                        choices=default_plays,
                        value=default_plays[0] if default_plays else None,
                        label="Play",
                    )
                    play_objective = gr.Textbox(
                        label="Plan objective",
                        placeholder="e.g., Turn post-storm claims into FORTIFIED upgrade conversions",
                        lines=2,
                    )
                    play_constraints = gr.Textbox(
                        label="Constraints and assumptions",
                        placeholder="e.g., limited grant support capacity, contractor onboarding lag",
                        lines=2,
                    )
                    with gr.Row():
                        horizon = gr.Dropdown(
                            choices=["30 days", "60 days", "90 days", "180 days"],
                            value="90 days",
                            label="Planning horizon",
                        )
                        budget_level = gr.Dropdown(
                            choices=["Lean", "Balanced", "Aggressive"],
                            value="Balanced",
                            label="Budget level",
                        )
                    play_search_online = gr.Checkbox(value=False, label="Search online (Tavily)")
                    play_online_focus = gr.Textbox(
                        label="Online search focus (optional)",
                        placeholder="e.g., best practices for contractor co-marketing in Louisiana",
                        visible=False,
                        value="",
                    )
                    play_temperature = gr.Slider(
                        minimum=0.0,
                        maximum=1.0,
                        value=0.6,
                        step=0.05,
                        label="Creativity",
                    )
                    generate_play = gr.Button("Generate Play Plan", variant="primary")

                with gr.Column(scale=5, elem_classes=["card"]):
                    # gr.Markdown("Play-based marketing plan will appear here.")
                    # play_sources = gr.Textbox(label="Source trace", lines=13)

                    play_output_m = gr.Markdown(label="Report", value="Play-based marketing plan will appear here.", padding=True)
                    play_output = gr.Code(label="Report", value="Play-based marketing plan will appear here.",
                                          language="markdown", lines=13, show_label=True, visible=False)
                    gr.Markdown("---")
                    play_sources = gr.Markdown(label="Source trace", value="Sources will appear here.", padding=True)
            state_single.change(update_play_options, inputs=state_single, outputs=play_selector)
            play_search_online.change(fn=toggle_online_focus, inputs=play_search_online, outputs=play_online_focus, show_progress="hidden")
            generate_play.click(
                fn=_generate_play_plan_output_stream,
                inputs=[
                    state_single,
                    play_selector,
                    play_objective,
                    play_constraints,
                    horizon,
                    budget_level,
                    play_search_online,
                    play_online_focus,
                    play_temperature,
                ],
                outputs=[play_output, play_output_m, play_sources, generate_play],
                show_progress="full",
            )



    return demo


STATES = load_state_catalog(CATALOG_PATH)
STATE_PLAYBOOKS = {state.slug: state.plays for state in STATES if state.plays}
STATES_BY_LABEL: Dict[str, StateInfo] = {state.label: state for state in STATES}
STATES_BY_SLUG: Dict[str, StateInfo] = {state.slug: state for state in STATES}
VECTOR_STORE, VECTOR_ERROR = load_vector_store()
QUESTION_BANK, QUESTION_BANK_ERROR = load_question_bank(QUESTIONS_PATH)


demo = build_app()


if __name__ == "__main__":
    demo.queue(default_concurrency_limit=10).launch(theme=gr.themes.Soft(), css=CUSTOM_CSS,)