app.py

"""
app.py — GraphReel Streamlit UI

Turns Google Drive files into a cohesive narrative + structured digest + explainer video
using a LangGraph pipeline backed by Vertex AI / Gemini.

Run with:  streamlit run app.py
"""

import queue as q_module
import threading
import time
import os
from pathlib import Path

import streamlit as st

# Reduce low-level gRPC C++ log noise in parallel pipeline runs.
os.environ.setdefault("GRPC_VERBOSITY", "ERROR")
os.environ.setdefault("GRPC_ENABLE_FORK_SUPPORT", "0")
os.environ.setdefault("GLOG_minloglevel", "2")

from drive import auth
from pipeline.graph import stream_graph, stream_video_graph

# ─────────────────────────────────────────────────────────────────────────────
# Page config
# ─────────────────────────────────────────────────────────────────────────────

st.set_page_config(
    page_title="GraphReel",
    page_icon="🎬",
    layout="centered",
)

DEFAULT_USER_INSTRUCTIONS = "Summarize in a high level explainer video."

# ─────────────────────────────────────────────────────────────────────────────
# Session state defaults
# ─────────────────────────────────────────────────────────────────────────────

if "result" not in st.session_state:
    st.session_state.result = None
if "warnings" not in st.session_state:
    st.session_state.warnings = []
if "steps" not in st.session_state:
    st.session_state.steps = []
if "video_path" not in st.session_state:
    st.session_state.video_path = None
if "video_script" not in st.session_state:
    st.session_state.video_script = None
if "video_title" not in st.session_state:
    st.session_state.video_title = ""
if "user_instructions" not in st.session_state:
    st.session_state.user_instructions = DEFAULT_USER_INSTRUCTIONS
if "instructions_input" not in st.session_state:
    st.session_state.instructions_input = st.session_state.user_instructions
if "url_input" not in st.session_state:
    st.session_state.url_input = ""
if "streaming_summaries" not in st.session_state:
    st.session_state.streaming_summaries = []
if "generating" not in st.session_state:
    st.session_state.generating = False
if "gen_status_label" not in st.session_state:
    st.session_state.gen_status_label = ""
if "gen_done" not in st.session_state:
    st.session_state.gen_done = False
if "gen_error" not in st.session_state:
    st.session_state.gen_error = None
if "gen_stopped" not in st.session_state:
    st.session_state.gen_stopped = False
# Queue and stop event are stored as non-serializable objects — Streamlit keeps them as-is
if "update_queue" not in st.session_state:
    st.session_state.update_queue = None
if "stop_event" not in st.session_state:
    st.session_state.stop_event = None


def _reset_page_state(clear_inputs: bool = False) -> None:
    """Reset generated outputs and pipeline state for a fresh run."""
    if st.session_state.stop_event:
        st.session_state.stop_event.set()

    st.session_state.result = None
    st.session_state.warnings = []
    st.session_state.steps = []
    st.session_state.video_path = None
    st.session_state.video_script = None
    st.session_state.video_title = ""
    st.session_state.streaming_summaries = []
    st.session_state.generating = False
    st.session_state.gen_status_label = ""
    st.session_state.gen_done = False
    st.session_state.gen_error = None
    st.session_state.gen_stopped = False
    st.session_state.update_queue = None
    st.session_state.stop_event = None

    if clear_inputs:
        st.session_state.user_instructions = DEFAULT_USER_INSTRUCTIONS
        st.session_state.instructions_input = DEFAULT_USER_INSTRUCTIONS
        st.session_state.url_input = ""

# ─────────────────────────────────────────────────────────────────────────────
# Background worker: runs both pipelines, pushes updates into a Queue
# ─────────────────────────────────────────────────────────────────────────────

def _run_pipelines(urls: list[str], stop_event: threading.Event, update_queue: q_module.Queue, user_instructions: str = ""):
    """Runs the digest then video pipeline; puts (type, ...) tuples onto the queue."""
    narrative = ""
    structured = ""
    summarize_count = [0]
    total_files = [0]
    scene_count = [0]
    total_scenes = [0]
    total_audio = [0.0]

    try:
        # ── Phase 1: Digest ──────────────────────────────────────────────────
        for node_name, node_state in stream_graph(urls, user_instructions):
            if stop_event.is_set():
                update_queue.put(("stopped",))
                return

            if node_name == "resolve_urls":
                n_drive = len(node_state.get("file_metas", []))
                n_web = len(node_state.get("web_contents", []))
                parts = []
                if n_drive:
                    parts.append(f"{n_drive} Drive file(s)")
                if n_web:
                    parts.append(f"{n_web} web page(s)")
                found = ", ".join(parts) or "0 sources"
                update_queue.put(("status", f"Found {found} — downloading contents..."))

            elif node_name == "fetch_files":
                total_files[0] = len(node_state.get("file_contents", []))
                update_queue.put(("status", f"Extracted {total_files[0]} source(s) — summarizing..."))

            elif node_name == "summarize_file":
                summarize_count[0] += 1
                name = ""
                summaries = node_state.get("file_summaries", [])
                if summaries:
                    name = summaries[-1].get("name", "")
                    update_queue.put(("new_summary", summaries[-1]))
                label = f"Summarizing file {summarize_count[0]}"
                if total_files[0]:
                    label += f" of {total_files[0]}"
                if name:
                    label += f": {name}"
                update_queue.put(("status", label + "..."))

            elif node_name == "synthesize":
                update_queue.put(("status", "Synthesizing final digest..."))

            elif node_name == "extract_topics":
                topics = node_state.get("research_topics", [])
                update_queue.put(("status",
                    f"Identified {len(topics)} research topic(s) — searching the web..."))

            elif node_name == "search_topic":
                research = node_state.get("web_research", [])
                topic = research[0]["topic"] if research else ""
                n_sources = len(research[0].get("sources", [])) if research else 0
                update_queue.put(("status",
                    f"Researched: \"{topic}\" ({n_sources} source(s) found)..."))

            elif node_name == "augment_narrative":
                update_queue.put(("status", "Weaving web research into narrative..."))

            elif node_name == "condense_narrative":
                update_queue.put(("status", "Condensing narrative for video..."))

            # Accumulate warnings and steps
            update_queue.put(("state", node_name, node_state))

            if node_name == "condense_narrative":
                narrative = node_state.get("narrative", narrative)
                structured = node_state.get("structured", structured)
            elif node_name == "augment_narrative":
                narrative = node_state.get("narrative", narrative)
                structured = node_state.get("structured", structured)
            elif node_name == "synthesize":
                narrative = node_state.get("narrative", "")
                structured = node_state.get("structured", "")

        if stop_event.is_set():
            update_queue.put(("stopped",))
            return

        update_queue.put(("status", "Digest complete — generating video script..."))

        # ── Phase 2: Video ───────────────────────────────────────────────────
        for node_name, node_state in stream_video_graph(narrative, structured):
            if stop_event.is_set():
                update_queue.put(("stopped",))
                return

            if node_name == "generate_script":
                scenes = node_state.get("script_scenes", [])
                title = node_state.get("video_title", "Explainer Video")
                total_scenes[0] = len(scenes)
                update_queue.put(("video_script", title, scenes))
                update_queue.put(("status", f"Script ready — {total_scenes[0]} scenes. Generating assets..."))

            elif node_name == "generate_scene_assets":
                scene_count[0] += 1
                assets = node_state.get("scene_assets", [])
                if assets:
                    total_audio[0] += assets[-1].get("audio_duration", 0)
                label = f"Generating assets for scene {scene_count[0]}"
                if total_scenes[0]:
                    label += f" of {total_scenes[0]}"
                update_queue.put(("status", label + "..."))
                if scene_count[0] == total_scenes[0]:
                    est_mins = max(1, int(total_audio[0] / 60) + 1)
                    update_queue.put(("status",
                        f"All assets ready — encoding {total_audio[0]:.0f}s of video "
                        f"(~{est_mins} min)..."))

            elif node_name == "assemble_video":
                video_path = node_state.get("video_path", "")
                update_queue.put(("video_path", video_path))
                update_queue.put(("status", f"Video written to {video_path}"))

            update_queue.put(("state", node_name, node_state))

        update_queue.put(("done",))

    except Exception as e:
        update_queue.put(("error", str(e)))


# ─────────────────────────────────────────────────────────────────────────────
# Header
# ─────────────────────────────────────────────────────────────────────────────

head_col, reset_col, drive_col = st.columns([5, 1, 2])
with head_col:
    st.title("🎬 GraphReel")
    st.caption(
        "Generate cohesive digests and AI explainer videos from web URLs and Google Drive sources."
    )

with reset_col:
    if st.session_state.video_path:
        if st.button("↺ Reset", help="Clear generated outputs and start a new video."):
            _reset_page_state(clear_inputs=True)
            st.rerun()

# Keep this state available for pipeline routing and UI hints
st.session_state.drive_connected = False
st.session_state.drive_user_email = ""

# ─────────────────────────────────────────────────────────────────────────────
# Google Drive auth check + compact top-right status
# ─────────────────────────────────────────────────────────────────────────────

try:
    creds = auth.get_credentials()
    user_email = auth.get_user_email(creds)
    st.session_state.drive_connected = True
    st.session_state.drive_user_email = user_email or ""
    with drive_col:
        if st.button("✅ Drive", help="Google Drive connected. Click to disconnect."):
            auth.revoke_credentials()
            st.rerun()

except auth.CredentialsFileNotFoundError as e:
    with drive_col:
        st.button("⚠️ Drive", disabled=True, help="Google OAuth credentials are missing.")
    st.warning(str(e))
    st.info(
        "To enable Google Drive links, create credentials.json per README instructions. "
        "Web URLs still work without Drive auth."
    )

except auth.NeedsAuthError:
    with drive_col:
        if st.button("🔌 Connect", type="primary", use_container_width=True, help="Connect Google Drive"):
            with st.spinner("Opening browser for Google sign-in..."):
                try:
                    auth.run_auth_flow()
                    st.rerun()
                except Exception as e:
                    st.error(f"Authentication failed: {e}")

st.divider()

# ─────────────────────────────────────────────────────────────────────────────
# URL input
# ─────────────────────────────────────────────────────────────────────────────

st.markdown("Instructions")
instructions_input = st.text_input(
    label="Instructions",
    key="instructions_input",
    placeholder="Summarize in a high level explainer video.",
    label_visibility="collapsed",
    help="Optional guidance for what to emphasize in the final explainer.",
    disabled=st.session_state.generating,
)
st.session_state.user_instructions = instructions_input

st.markdown(
    "Paste GDrive or web URLs, one per line"
)

url_input = st.text_area(
    label="URLs",
    key="url_input",
    placeholder=(
        "https://drive.google.com/drive/folders/...\n"
        "https://docs.google.com/document/d/.../edit\n"
        "https://example.com/blog/some-article"
    ),
    height=120,
    label_visibility="collapsed",
    disabled=st.session_state.generating,
)

# ─────────────────────────────────────────────────────────────────────────────
# Generate / Stop buttons
# ─────────────────────────────────────────────────────────────────────────────

if st.session_state.generating:
    # Make the Stop button red — scoped to the block that contains a disabled button
    st.markdown(
        """
        <style>
        div[data-testid="stHorizontalBlock"]:has(button:disabled)
            [data-testid="stBaseButton-secondary"] {
            background-color: #dc2626 !important;
            border-color: #dc2626 !important;
            color: white !important;
        }
        div[data-testid="stHorizontalBlock"]:has(button:disabled)
            [data-testid="stBaseButton-secondary"]:hover {
            background-color: #b91c1c !important;
            border-color: #b91c1c !important;
        }
        </style>
        """,
        unsafe_allow_html=True,
    )
    btn_col, stop_col = st.columns([4, 1])
    with btn_col:
        st.button(
            "Generating…",
            type="primary",
            use_container_width=True,
            disabled=True,
        )
    with stop_col:
        if st.button("Stop", type="secondary", use_container_width=True):
            if st.session_state.stop_event:
                st.session_state.stop_event.set()
            st.session_state.generating = False
            st.session_state.gen_stopped = True
            st.rerun()
else:
    generate_clicked = st.button(
        "Generate Digest & Video",
        type="primary",
        use_container_width=True,
    )

    if generate_clicked:
        urls = [u.strip() for u in url_input.strip().splitlines() if u.strip()]
        if not urls:
            st.warning("Please paste at least one URL.")
        else:
            _reset_page_state(clear_inputs=False)
            st.session_state.gen_status_label = "Starting pipeline..."

            stop_event = threading.Event()
            update_queue = q_module.Queue()
            st.session_state.stop_event = stop_event
            st.session_state.update_queue = update_queue
            st.session_state.generating = True

            t = threading.Thread(
                target=_run_pipelines,
                args=(urls, stop_event, update_queue, st.session_state.user_instructions),
                daemon=True,
            )
            t.start()
            st.rerun()

# ─────────────────────────────────────────────────────────────────────────────
# Poll the update queue while generating
# ─────────────────────────────────────────────────────────────────────────────

if st.session_state.generating and st.session_state.update_queue is not None:
    uq: q_module.Queue = st.session_state.update_queue

    # Drain all available messages
    while True:
        try:
            msg = uq.get_nowait()
        except q_module.Empty:
            break

        kind = msg[0]

        if kind == "status":
            st.session_state.gen_status_label = msg[1]

        elif kind == "state":
            _, node_name, node_state = msg
            for w in node_state.get("warnings", []):
                if w not in st.session_state.warnings:
                    st.session_state.warnings.append(w)
            for s in node_state.get("steps", []):
                if s not in st.session_state.steps:
                    st.session_state.steps.append(s)
            if node_name in ("synthesize", "augment_narrative", "condense_narrative"):
                st.session_state.result = node_state

        elif kind == "new_summary":
            summary = msg[1]
            existing_ids = {s["file_id"] for s in st.session_state.streaming_summaries}
            if summary["file_id"] not in existing_ids:
                st.session_state.streaming_summaries.append(summary)

        elif kind == "video_script":
            st.session_state.video_title = msg[1]
            st.session_state.video_script = msg[2]

        elif kind == "video_path":
            st.session_state.video_path = msg[1]

        elif kind == "done":
            st.session_state.generating = False
            st.session_state.gen_done = True
            st.session_state.gen_status_label = "Complete!"

        elif kind == "stopped":
            st.session_state.generating = False
            st.session_state.gen_stopped = True

        elif kind == "error":
            st.session_state.generating = False
            st.session_state.gen_error = msg[1]

    # If still running, show animated progress and keep polling
    if st.session_state.generating:
        st.markdown(
            f"""
            <style>
            @keyframes _gr_spin {{
                to {{ transform: rotate(360deg); }}
            }}
            </style>
            <div style="display:flex; align-items:center; gap:12px; padding:14px 16px;
                        background:#e8f4fb; border-radius:6px; border-left:4px solid #1f77b4;">
                <div style="width:18px; height:18px; flex-shrink:0;
                            border:3px solid #1f77b4; border-top-color:transparent;
                            border-radius:50%;
                            animation:_gr_spin 0.75s linear infinite;"></div>
                <span style="color:#1f77b4; font-size:0.95rem; font-weight:500;">
                    {st.session_state.gen_status_label}
                </span>
            </div>
            """,
            unsafe_allow_html=True,
        )
        # Show script preview as soon as it's available (while assets are still generating)
        if st.session_state.video_script:
            st.markdown("<div style='height: 8px;'></div>", unsafe_allow_html=True)
            with st.expander(
                f"Video script preview ({len(st.session_state.video_script)} scenes) — generating assets…",
                expanded=True,
            ):
                for scene in st.session_state.video_script:
                    st.markdown(f"**Scene {scene['scene_number']}**")
                    st.markdown(f"*Narration:* {scene['narration']}")
                    st.markdown(f"*Overlay:* {scene['overlay_text']}")
                    st.markdown(f"*Image prompt:* {scene['image_prompt']}")
                    st.markdown("---")

        time.sleep(0.4)
        st.rerun()

# ─────────────────────────────────────────────────────────────────────────────
# Terminal status messages
# ─────────────────────────────────────────────────────────────────────────────

if st.session_state.gen_done:
    st.success(f"✓ {st.session_state.gen_status_label}")

if st.session_state.gen_stopped:
    st.warning("Generation stopped.")

if st.session_state.gen_error:
    st.error(f"Pipeline failed: {st.session_state.gen_error}")

elif st.session_state.gen_status_label and not st.session_state.generating and not st.session_state.gen_done and not st.session_state.gen_stopped:
    # Show last status if there was one (e.g. after a rerun)
    pass

# ─────────────────────────────────────────────────────────────────────────────
# Streaming per-file summaries (shown as each file completes, even mid-generation)
# ─────────────────────────────────────────────────────────────────────────────

if st.session_state.streaming_summaries:
    label = f"Per-file summaries ({len(st.session_state.streaming_summaries)} complete)"
    if st.session_state.generating:
        label += " — in progress…"
    with st.expander(label, expanded=st.session_state.generating):
        for fs in st.session_state.streaming_summaries:
            st.markdown(f"**{fs['name']}**")
            st.markdown(fs["summary"])
            st.markdown("---")

# ─────────────────────────────────────────────────────────────────────────────
# Warnings banner
# ─────────────────────────────────────────────────────────────────────────────

if st.session_state.warnings:
    with st.expander(f"⚠ {len(st.session_state.warnings)} warning(s)", expanded=False):
        for w in st.session_state.warnings:
            st.markdown(f"- {w}")

# ─────────────────────────────────────────────────────────────────────────────
# Digest output
# ─────────────────────────────────────────────────────────────────────────────

if st.session_state.result:
    result = st.session_state.result
    narrative = result.get("narrative", "")
    structured = result.get("structured", "")

    st.divider()
    st.subheader("Digest")

    tab_narrative, tab_structured = st.tabs(["Narrative Briefing", "Structured Digest"])

    with tab_narrative:
        st.markdown(
            "_Flowing prose optimized for reading aloud — your future video script._"
        )
        if narrative:
            st.markdown(narrative)
        else:
            st.info("Narrative not yet generated. Try running the pipeline again.")
        if narrative:
            st.download_button(
                label="Download as .txt (video script draft)",
                data=narrative,
                file_name="graphreel_narrative.txt",
                mime="text/plain",
            )

    with tab_structured:
        st.markdown(
            "_Organized by theme for quick scanning and reference._"
        )
        if structured:
            st.markdown(structured)
        else:
            st.info("Structured digest not yet generated. Try running the pipeline again.")

    with st.expander("Debug — pipeline steps", expanded=False):
        for step in st.session_state.steps:
            st.markdown(f"- {step}")

# ─────────────────────────────────────────────────────────────────────────────
# Video output
# ─────────────────────────────────────────────────────────────────────────────

if st.session_state.video_path:
    st.divider()
    st.subheader("Explainer Video")
    st.video(st.session_state.video_path)

    video_filename = Path(st.session_state.video_path).name
    with open(st.session_state.video_path, "rb") as vf:
        st.download_button(
            label="Download MP4",
            data=vf,
            file_name=video_filename,
            mime="video/mp4",
        )

    with st.expander("Upload to YouTube", expanded=False):
        st.markdown(
            "_YouTube doesn't support URL-based pre-fill, so use the fields below "
            "to copy your metadata, then drag the MP4 into YouTube Studio._"
        )
        st.link_button(
            "Open YouTube Studio ↗",
            url="https://studio.youtube.com/channel/upload",
            type="primary",
            use_container_width=True,
        )
        yt_title = st.text_input(
            "Title",
            value=st.session_state.video_title,
            help="Copy this into the YouTube title field",
        )
        narrative_for_yt = ""
        if st.session_state.result:
            narrative_for_yt = st.session_state.result.get("narrative", "")[:4900]
        st.text_area(
            "Description",
            value=narrative_for_yt,
            height=200,
            help="Copy this into the YouTube description field (trimmed to 4900 chars)",
        )
        st.markdown("**File location:**")
        st.code(st.session_state.video_path, language=None)

if st.session_state.video_script:
    with st.expander(
        f"Scene-by-scene script ({len(st.session_state.video_script)} scenes)",
        expanded=False,
    ):
        for scene in st.session_state.video_script:
            st.markdown(f"**Scene {scene['scene_number']}**")
            st.markdown(f"*Narration:* {scene['narration']}")
            st.markdown(f"*Overlay:* {scene['overlay_text']}")
            st.markdown(f"*Image prompt:* {scene['image_prompt']}")
            st.markdown("---")