app.py

"""NotebookLM — AI-Powered Study Companion (Gradio)."""

import gradio as gr

from state import (
    UserData,
    create_default_user_data,
    create_notebook,
    delete_notebook,
    rename_notebook,
    get_active_notebook,
    get_notebook_choices,
    get_latest_artifact,
)
from persistence.storage_service import StorageService
from theme import dark_theme, CUSTOM_CSS, SIDEBAR_LOGO_HTML, WELCOME_HTML, NO_NOTEBOOKS_HTML
from pages.chat import (
    format_chatbot_messages,
    render_no_sources_warning,
    handle_chat_submit,
    handle_clear_chat,
)
from pages.sources import (
    render_source_header,
    render_source_list,
    get_source_choices,
    handle_file_upload,
    handle_url_add,
    handle_source_delete,
    run_ingestion_pipeline,
)
from pages.artifacts import (
    render_no_sources_gate,
    has_sources,
    render_conv_summary_section,
    handle_gen_conv_summary,
    render_doc_summary_section,
    handle_gen_doc_summary,
    render_podcast_section,
    handle_gen_podcast,
    render_quiz_section,
    handle_gen_quiz,
    has_any_summary,
    download_artifact,
    _save_artifact_to_temp_file,
)

# ── Helpers ──────────────────────────────────────────────────────────────────

def render_notebook_header(state: UserData) -> str:
    nb = get_active_notebook(state)
    if not nb:
        return NO_NOTEBOOKS_HTML
    src = len(nb.sources)
    msg = len(nb.messages)
    art = len(nb.artifacts)
    return (
        f'<div class="notebook-header">'
        f'<h2>{nb.title}</h2>'
        f'<div class="meta">{src} sources &bull; {msg} messages &bull; {art} artifacts</div>'
        f'</div>'
    )


def render_user_info(state: UserData) -> str:
    if not state:
        return ""
    return (
        f'<p style="font-size:0.82rem; color:#707088; margin:4px 0 0 0;">'
        f'Signed in as <strong style="color:#a0a0f0;">{state.user_name}</strong></p>'
    )


# ── Download Handlers ────────────────────────────────────────────────────────

def download_conversation_summary(state: UserData) -> str | None:
    """Download the latest conversation summary as markdown."""
    nb = get_active_notebook(state)
    if not nb:
        return None
    artifact = get_latest_artifact(nb, "conversation_summary")
    if not artifact:
        return None
    return _save_artifact_to_temp_file(artifact)


def download_document_summary(state: UserData) -> str | None:
    """Download the latest document summary as markdown."""
    nb = get_active_notebook(state)
    if not nb:
        return None
    artifact = get_latest_artifact(nb, "document_summary")
    if not artifact:
        return None
    return _save_artifact_to_temp_file(artifact)


def download_podcast(state: UserData) -> str | None:
    """Download the latest podcast audio."""
    nb = get_active_notebook(state)
    if not nb:
        return None
    artifact = get_latest_artifact(nb, "podcast")
    if not artifact:
        return None
    return _save_artifact_to_temp_file(artifact)


def download_quiz(state: UserData) -> str | None:
    """Download the latest quiz as markdown."""
    nb = get_active_notebook(state)
    if not nb:
        return None
    artifact = get_latest_artifact(nb, "quiz")
    if not artifact:
        return None
    return _save_artifact_to_temp_file(artifact)


def refresh_all(state: UserData):
    """Refresh all display components after state change. Returns a tuple of all outputs."""
    nb = get_active_notebook(state)
    choices = get_notebook_choices(state) if state else []
    active_id = state.active_notebook_id if state else None

    has_nb = nb is not None
    has_src = has_nb and len(nb.sources) > 0

    ready_sources = [s.filename for s in nb.sources if s.status == "ready"] if has_nb else []

    return (
        state,
        # Sidebar
        gr.update(choices=choices, value=active_id),
        render_user_info(state),
        # Header
        render_notebook_header(state),
        # Chat tab
        format_chatbot_messages(state) if has_nb else [],
        render_no_sources_warning(state),
        # Sources tab
        render_source_header(state),
        render_source_list(state),
        gr.update(choices=get_source_choices(state)),
        # Artifacts tab
        gr.update(value=render_no_sources_gate(state), visible=not has_src),
        gr.update(visible=has_src),   # artifacts_content visible
        # Artifact sub-sections
        render_conv_summary_section(state),
        gr.update(choices=ready_sources, value=ready_sources),
        render_doc_summary_section(state),
        render_podcast_section(state),
        render_quiz_section(state),
    )


MARK_DIRTY_JS = "() => { window.__unsaved = true; }"
MARK_CLEAN_JS = "() => { window.__unsaved = false; }"

BEFOREUNLOAD_JS = """
() => {
    window.addEventListener('beforeunload', function(e) {
        if (window.__unsaved) {
            e.preventDefault();
            e.returnValue = '';
        }
    });
}
"""


# ── Build the App ────────────────────────────────────────────────────────────

with gr.Blocks(css=CUSTOM_CSS, theme=dark_theme, title="NotebookLM", js=BEFOREUNLOAD_JS) as demo:

    user_state = gr.State(value=None)

    # ══ Auth Gate ════════════════════════════════════════════════════════════
    with gr.Column(visible=True, elem_id="auth-gate") as auth_gate:
        gr.HTML(WELCOME_HTML)
        gr.LoginButton(elem_id="login-btn")

    # ══ Main App (hidden until login) ════════════════════════════════════════
    with gr.Row(visible=False) as main_app:

        # ── Sidebar ──────────────────────────────────────────────────────────
        with gr.Column(scale=1, min_width=280, elem_id="sidebar"):
            gr.HTML(SIDEBAR_LOGO_HTML)
            user_info_html = gr.HTML("")

            gr.Markdown("---")

            # Create notebook
            new_nb_name = gr.Textbox(
                placeholder="e.g. Biology 101",
                show_label=False,
                container=False,
            )
            create_nb_btn = gr.Button("+ New Notebook", variant="primary", size="sm")

            gr.HTML('<div style="height:8px;"></div>')

            # Notebook selector
            notebook_selector = gr.Radio(
                choices=[],
                label="Notebooks",
                elem_id="notebook-selector",
            )

            gr.Markdown("---")

            # Rename
            rename_input = gr.Textbox(
                placeholder="New name...",
                show_label=False,
                container=False,
            )
            rename_btn = gr.Button("Rename", size="sm")

            # Delete
            delete_btn = gr.Button("Delete Notebook", variant="stop", size="sm")

            gr.Markdown("---")

            # Save
            save_btn = gr.Button("Save", variant="primary", size="sm", elem_id="save-btn")
            save_status_html = gr.HTML("")

            gr.HTML(
                '<p style="font-size:0.75rem; color:#50506a; text-align:center; margin-top:16px;">'
                'Built with Gradio on HF Spaces</p>'
            )

        # ── Main Content ─────────────────────────────────────────────────────
        with gr.Column(scale=4, elem_id="main-content"):

            notebook_header = gr.HTML(NO_NOTEBOOKS_HTML)

            with gr.Tabs(elem_id="main-tabs") as main_tabs:

                # ── Chat Tab ─────────────────────────────────────────────────
                with gr.TabItem("Chat", id=0):
                    chat_warning = gr.HTML("")
                    chatbot = gr.Chatbot(
                        value=[],
                        type="messages",
                        height=480,
                        elem_id="chatbot",
                        show_label=False,
                    )
                    with gr.Row():
                        chat_input = gr.Textbox(
                            placeholder="Ask a question about your sources...",
                            show_label=False,
                            container=False,
                            scale=5,
                        )
                        clear_chat_btn = gr.Button("Clear", scale=1)

                # ── Sources Tab ──────────────────────────────────────────────
                with gr.TabItem("Sources", id=1):
                    source_header = gr.HTML("")

                    with gr.Row():
                        with gr.Column():
                            gr.HTML(
                                '<p style="font-weight:600; font-size:0.9rem; color:#b0b0c8; margin-bottom:8px;">'
                                'Upload Files</p>'
                            )
                            file_uploader = gr.File(
                                file_count="multiple",
                                file_types=[".pdf", ".pptx", ".txt"],
                                label="Drop files here",
                                show_label=False,
                            )
                        with gr.Column():
                            gr.HTML(
                                '<p style="font-weight:600; font-size:0.9rem; color:#b0b0c8; margin-bottom:8px;">'
                                'Add Web Source</p>'
                            )
                            url_input = gr.Textbox(
                                placeholder="https://example.com  or  YouTube link",
                                show_label=False,
                                container=False,
                            )
                            add_url_btn = gr.Button("Add URL", variant="primary")

                    gr.Markdown("---")
                    source_list_html = gr.HTML("")

                    with gr.Row():
                        source_selector = gr.Dropdown(
                            choices=[],
                            label="Select source to delete",
                            scale=3,
                            elem_id="source-delete-dropdown",
                            allow_custom_value=True,
                        )
                        delete_source_btn = gr.Button("Delete Source", variant="stop", scale=1)

                # ── Artifacts Tab ────────────────────────────────────────────
                with gr.TabItem("Artifacts", id=2):

                    # No-sources gate
                    no_sources_msg = gr.HTML("", visible=True)

                    with gr.Column(visible=False) as artifacts_content:

                        with gr.Tabs(elem_id="artifact-tabs"):

                            # Summary sub-tab
                            with gr.TabItem("Summary"):
                                # Conversation Summary
                                gr.HTML(
                                    '<div class="artifact-section-header">'
                                    '<div class="artifact-section-icon" style="background:rgba(102,126,234,0.12);">💬</div>'
                                    '<div><span style="font-weight:600; font-size:1rem; color:#e0e0f0;">Conversation Summary</span>'
                                    '<p style="font-size:0.82rem; color:#808098; margin:2px 0 0 0;">'
                                    'Summarize your chat history.</p></div></div>'
                                )
                                with gr.Row():
                                    conv_style_radio = gr.Radio(
                                        choices=["brief", "detailed"],
                                        value="detailed",
                                        label="Style",
                                        scale=2,
                                    )
                                    gen_conv_sum_btn = gr.Button(
                                        "Generate Conversation Summary",
                                        variant="primary",
                                        scale=2,
                                    )
                                conv_summary_html = gr.Markdown("")
                                download_conv_btn = gr.DownloadButton(
                                    label="📥 Download Summary (.md)",
                                    variant="secondary",
                                    scale=1,
                                )

                                gr.HTML('<div style="margin:30px 0; border-top:1px solid rgba(255,255,255,0.06);"></div>')

                                # Document Summary
                                gr.HTML(
                                    '<div class="artifact-section-header">'
                                    '<div class="artifact-section-icon" style="background:rgba(34,197,94,0.12);">📄</div>'
                                    '<div><span style="font-weight:600; font-size:1rem; color:#e0e0f0;">Document Summary</span>'
                                    '<p style="font-size:0.82rem; color:#808098; margin:2px 0 0 0;">'
                                    'Summarize content from your uploaded sources.</p></div></div>'
                                )
                                doc_source_selector = gr.CheckboxGroup(
                                    choices=[],
                                    label="Select sources to summarize",
                                    value=[],
                                )
                                with gr.Row():
                                    doc_style_radio = gr.Radio(
                                        choices=["brief", "detailed"],
                                        value="detailed",
                                        label="Style",
                                        scale=2,
                                    )
                                    gen_doc_sum_btn = gr.Button(
                                        "Generate Document Summary",
                                        variant="primary",
                                        scale=2,
                                    )
                                doc_summary_html = gr.Markdown("")
                                download_doc_btn = gr.DownloadButton(
                                    label="📥 Download Summary (.md)",
                                    variant="secondary",
                                    scale=1,
                                )

                            # Podcast sub-tab
                            with gr.TabItem("Podcast"):
                                gr.HTML(
                                    '<div style="margin-bottom:20px;">'
                                    '<span style="font-weight:600; font-size:1rem; color:#e0e0f0;">Generate Podcast</span>'
                                    '<p style="font-size:0.85rem; color:#808098; margin-top:4px;">'
                                    'Create a conversational podcast episode from your summary.</p></div>'
                                )
                                gen_podcast_btn = gr.Button("Generate Podcast", variant="primary")
                                podcast_html = gr.Markdown("")
                                download_podcast_btn = gr.DownloadButton(
                                    label="🎵 Download Audio (.mp3)",
                                    variant="secondary",
                                )

                            # Quiz sub-tab
                            with gr.TabItem("Quiz"):
                                gr.HTML(
                                    '<div style="margin-bottom:20px;">'
                                    '<span style="font-weight:600; font-size:1rem; color:#e0e0f0;">Generate Quiz</span>'
                                    '<p style="font-size:0.85rem; color:#808098; margin-top:4px;">'
                                    'Create multiple-choice questions from your sources.</p></div>'
                                )
                                with gr.Row():
                                    quiz_num_radio = gr.Radio(
                                        choices=[5, 10],
                                        value=5,
                                        label="Number of questions",
                                        scale=2,
                                    )
                                    gen_quiz_btn = gr.Button(
                                        "Generate Quiz",
                                        variant="primary",
                                        scale=2,
                                    )
                                quiz_html = gr.Markdown("")
                                download_quiz_btn = gr.DownloadButton(
                                    label="📋 Download Quiz (.html)",
                                    variant="secondary",
                                    scale=1,
                                )

    # ── All refresh outputs (must match refresh_all return order) ─────────
    refresh_outputs = [
        user_state,
        notebook_selector,
        user_info_html,
        notebook_header,
        chatbot,
        chat_warning,
        source_header,
        source_list_html,
        source_selector,
        no_sources_msg,
        artifacts_content,
        conv_summary_html,
        doc_source_selector,
        doc_summary_html,
        podcast_html,
        quiz_html,
    ]

    # ══ Event Handlers ═══════════════════════════════════════════════════════

    # ── Auth: on page load ───────────────────────────────────────────────────
    def on_app_load(profile: gr.OAuthProfile | None):
        print(f"[on_app_load] profile={profile}", flush=True)
        if profile is None:
            return None, gr.update(visible=True), gr.update(visible=False)
        print(f"[on_app_load] Loading data for user: {profile.username}", flush=True)
        state = StorageService.load_user_data(profile.username, profile.name)
        if state is None:
            print(f"[on_app_load] No saved data found — creating default", flush=True)
            state = create_default_user_data(profile.username, profile.name)
        else:
            print(f"[on_app_load] Loaded {len(state.notebooks)} notebook(s)", flush=True)
        return state, gr.update(visible=False), gr.update(visible=True)

    demo.load(
        fn=on_app_load,
        inputs=None,
        outputs=[user_state, auth_gate, main_app],
        api_name=False,
    ).then(
        fn=refresh_all,
        inputs=[user_state],
        outputs=refresh_outputs,
        api_name=False,
    )

    # ── Sidebar: Create notebook ─────────────────────────────────────────────
    def handle_create_notebook(name, state):
        if not name or not name.strip() or not state:
            return (state,) + refresh_all(state)[1:] + ("",)
        state = create_notebook(state, name.strip())
        return (state,) + refresh_all(state)[1:] + ("",)

    create_nb_btn.click(
        fn=handle_create_notebook,
        inputs=[new_nb_name, user_state],
        outputs=refresh_outputs + [new_nb_name],
        api_name=False,
    ).then(fn=None, js=MARK_DIRTY_JS)

    # ── Sidebar: Select notebook ─────────────────────────────────────────────
    def handle_select_notebook(nb_id, state):
        if not state or not nb_id:
            return refresh_all(state)
        state.active_notebook_id = nb_id
        return refresh_all(state)

    notebook_selector.change(
        fn=handle_select_notebook,
        inputs=[notebook_selector, user_state],
        outputs=refresh_outputs,
        api_name=False,
    )

    # ── Sidebar: Delete notebook ─────────────────────────────────────────────
    def handle_delete_notebook(state):
        if not state or not state.active_notebook_id:
            return refresh_all(state)
        state = delete_notebook(state, state.active_notebook_id)
        return refresh_all(state)

    delete_btn.click(
        fn=handle_delete_notebook,
        inputs=[user_state],
        outputs=refresh_outputs,
        api_name=False,
    ).then(fn=None, js=MARK_DIRTY_JS)

    # ── Sidebar: Rename notebook ─────────────────────────────────────────────
    def handle_rename_notebook(new_name, state):
        if not state or not state.active_notebook_id or not new_name or not new_name.strip():
            return refresh_all(state)
        state = rename_notebook(state, state.active_notebook_id, new_name.strip())
        return refresh_all(state)

    rename_btn.click(
        fn=handle_rename_notebook,
        inputs=[rename_input, user_state],
        outputs=refresh_outputs,
        api_name=False,
    ).then(fn=None, js=MARK_DIRTY_JS)

    # ── Sidebar: Save ────────────────────────────────────────────────────────
    def handle_save(state):
        print(f"[handle_save] Called, state={state is not None}", flush=True)
        if not state:
            return '<p style="color:#ef4444; font-size:0.8rem;">No data to save.</p>'
        try:
            StorageService.save_user_data(state)
            return '<p style="color:#22c55e; font-size:0.8rem;">Saved successfully!</p>'
        except Exception as e:
            print(f"[handle_save] Error: {e}", flush=True)
            return f'<p style="color:#ef4444; font-size:0.8rem;">Save failed: {e}</p>'

    save_btn.click(
        fn=handle_save,
        inputs=[user_state],
        outputs=[save_status_html],
        api_name=False,
    ).then(fn=None, js=MARK_CLEAN_JS)

    # ── Chat: Submit message ─────────────────────────────────────────────────
    chat_input.submit(
        fn=handle_chat_submit,
        inputs=[chat_input, user_state],
        outputs=[user_state, chatbot, chat_input, chat_warning],
        api_name=False,
    )

    # ── Chat: Clear ──────────────────────────────────────────────────────────
    clear_chat_btn.click(
        fn=handle_clear_chat,
        inputs=[user_state],
        outputs=[user_state, chatbot, chat_warning],
        api_name=False,
    )

    # ── Sources: File upload ─────────────────────────────────────────────────
    def handle_upload_and_ingest(files, state):
        state, _, _, _ = handle_file_upload(files, state)
        yield refresh_all(state)          # UI updates immediately — sources show "Processing..."
        run_ingestion_pipeline(state)     # Slow: extract, chunk, embed
        yield refresh_all(state)          # UI updates again — sources show "Ready"

    file_uploader.upload(
        fn=handle_upload_and_ingest,
        inputs=[file_uploader, user_state],
        outputs=refresh_outputs,
        api_name=False,
    ).then(fn=None, js=MARK_DIRTY_JS)

    # ── Sources: Add URL ─────────────────────────────────────────────────────
    def handle_url_add_and_ingest(url, state):
        state, _, _, _, _ = handle_url_add(url, state)
        yield refresh_all(state) + ("",)  # UI updates immediately + clear URL input
        run_ingestion_pipeline(state)     # Slow: extract, chunk, embed
        yield refresh_all(state) + ("",)  # UI updates again — source shows "Ready"

    add_url_btn.click(
        fn=handle_url_add_and_ingest,
        inputs=[url_input, user_state],
        outputs=refresh_outputs + [url_input],
        api_name=False,
    ).then(fn=None, js=MARK_DIRTY_JS)

    # ── Sources: Delete source ───────────────────────────────────────────────
    delete_source_btn.click(
        fn=handle_source_delete,
        inputs=[source_selector, user_state],
        outputs=[user_state, source_list_html, source_header, source_selector],
        api_name=False,
    ).then(
        fn=refresh_all,
        inputs=[user_state],
        outputs=refresh_outputs,
        api_name=False,
    ).then(fn=None, js=MARK_DIRTY_JS)

    # ── Artifacts: Conversation summary ──────────────────────────────────────
    gen_conv_sum_btn.click(
        fn=handle_gen_conv_summary,
        inputs=[conv_style_radio, user_state],
        outputs=[user_state, conv_summary_html],
        api_name=False,
    )

    # ── Artifacts: Document summary ──────────────────────────────────────────
    gen_doc_sum_btn.click(
        fn=handle_gen_doc_summary,
        inputs=[doc_style_radio, doc_source_selector, user_state],
        outputs=[user_state, doc_summary_html],
        api_name=False,
    )

    # ── Artifacts: Podcast ───────────────────────────────────────────────────
    gen_podcast_btn.click(
        fn=handle_gen_podcast,
        inputs=[user_state],
        outputs=[user_state, podcast_html],
        api_name=False,
    )

    # ── Artifacts: Quiz ──────────────────────────────────────────────────────
    gen_quiz_btn.click(
        fn=handle_gen_quiz,
        inputs=[quiz_num_radio, user_state],
        outputs=[user_state, quiz_html],
        api_name=False,
    )

    # ── Artifacts: Download buttons ──────────────────────────────────────────
    download_conv_btn.click(
        fn=download_conversation_summary,
        inputs=[user_state],
        outputs=[download_conv_btn],
        api_name=False,
    )

    download_doc_btn.click(
        fn=download_document_summary,
        inputs=[user_state],
        outputs=[download_doc_btn],
        api_name=False,
    )

    download_podcast_btn.click(
        fn=download_podcast,
        inputs=[user_state],
        outputs=[download_podcast_btn],
        api_name=False,
    )

    download_quiz_btn.click(
        fn=download_quiz,
        inputs=[user_state],
        outputs=[download_quiz_btn],
        api_name=False,
    )


# ── Launch ───────────────────────────────────────────────────────────────────
if __name__ == "__main__":
    demo.launch(show_api=False, ssr_mode=False, allowed_paths=["assets/podcasts"])