#!/usr/bin/env python3
"""
AI Story Livestream — RTMP livestream example for the Reactor Python SDK.

Connects to the livecore model (LongLive) and generates an endless stream of
AI-generated visual stories.  Each story is written by OpenAI ChatGPT as a
sequence of 7 cinematic scene descriptions.  The scenes are scheduled as prompt
transitions across the 240-frame generation window (6 transitions), producing
smooth visual storytelling streamed live to YouTube (or any RTMP endpoint).

While the current story is being rendered, the **next** story is generated in
the background (pipeline), so there is no dead time between stories.

If the WebRTC connection to Reactor drops, the app automatically reconnects and
starts a fresh session — the RTMP stream to YouTube stays alive.

Requirements:
    - ffmpeg must be installed and available in PATH
    - OPENAI_API_KEY environment variable must be set
    - A Reactor API key (or --local for local development)

Usage:
    python main.py --api-key KEY --rtmp-url rtmp://a.rtmp.youtube.com/live2/STREAM_KEY
    python main.py --local --rtmp-url rtmp://a.rtmp.youtube.com/live2/STREAM_KEY
    python main.py --local --rtmp-url rtmp://... --audio background.mp3
"""

from __future__ import annotations

import argparse
import asyncio
import json
import logging
import os
import random
import shutil
import subprocess
import sys
import warnings
from pathlib import Path

import aiohttp

import numpy as np
from numpy.typing import NDArray

# Add parent directories to path for development
sys.path.insert(0, str(__file__).rsplit("/", 3)[0] + "/src")

from reactor_sdk import MessageScope, Reactor, ReactorStatus
from reactor_sdk.types import ReactorError

# =============================================================================
# OpenAI client (lazy import so we get a clear error message)
# =============================================================================

try:
    from openai import AsyncOpenAI
except ImportError:
    print(
        "The 'openai' package is required.  Install it with:\n"
        "  pip install openai"
    )
    sys.exit(1)

# =============================================================================
# Logging
# =============================================================================

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s",
)
logger = logging.getLogger(__name__)

# Suppress noisy third-party warnings
logging.getLogger("aiortc.codecs.vpx").setLevel(logging.ERROR)
logging.getLogger("aiortc.codecs.h264").setLevel(logging.ERROR)
logging.getLogger("aioice.ice").setLevel(logging.WARNING)
warnings.filterwarnings("ignore", category=DeprecationWarning)

# =============================================================================
# Constants
# =============================================================================

# Livecore model output resolution (from its config: height=480, width=832).
# These MUST match the model — a mismatch causes garbled "TV static" in the
# RTMP stream because ffmpeg interprets raw bytes with the wrong row stride.
MODEL_WIDTH = 832
MODEL_HEIGHT = 480

# Livecore generates 240 frames per generation cycle.  We schedule 7 prompts
# (the opening scene + 6 transitions) spread evenly across the window.
TOTAL_FRAMES = 240
SCENE_FRAMES = [0, 34, 68, 102, 136, 170, 204]
NUM_SCENES = len(SCENE_FRAMES)  # 7

# Reconnection
RECONNECT_DELAY_S = 5
MAX_RECONNECT_ATTEMPTS = 100  # effectively unlimited

# =============================================================================
# Parse CLI arguments
# =============================================================================

parser = argparse.ArgumentParser(
    description="AI Story Livestream — endless AI-generated stories via RTMP"
)
parser.add_argument(
    "--model", "-m", default="livecore",
    help="Model name (default: livecore)",
)
parser.add_argument("--api-key", "-k", help="Reactor API key")
parser.add_argument(
    "--local", "-l", action="store_true", help="Use local coordinator"
)
parser.add_argument(
    "--rtmp-url", "-r", required=True,
    help="RTMP URL including stream key",
)
parser.add_argument(
    "--fps", "-f", type=int, default=20,
    help="Output frame rate (default: 20)",
)
parser.add_argument(
    "--audio", "-a",
    help="Path to MP3 file to loop as background audio",
)
parser.add_argument(
    "--topics-file",
    default=os.path.join(os.path.dirname(__file__), "topics.txt"),
    help="Path to topics file (default: ./topics.txt)",
)
parser.add_argument(
    "--openai-model",
    default="gpt-4o-mini",
    help="OpenAI model for story generation (default: gpt-4o-mini)",
)
parser.add_argument(
    "--youtube-video-id",
    help="YouTube video/stream ID to read live chat from (e.g. dQw4w9WgXcQ)",
)
parser.add_argument(
    "--chat-messages",
    type=int,
    default=1,
    help="Minimum new !p messages required to override topics (default: 1)",
)
parser.add_argument(
    "--verbose", "-v", action="store_true", help="Enable debug logging"
)
args = parser.parse_args()

if args.verbose:
    logging.getLogger().setLevel(logging.DEBUG)

if not args.local and not args.api_key:
    logger.error("Either --local or --api-key must be provided")
    sys.exit(1)

if shutil.which("ffmpeg") is None:
    logger.error(
        "ffmpeg not found! Please install ffmpeg and ensure it's in your PATH.\n"
        "  macOS:   brew install ffmpeg\n"
        "  Linux:   apt install ffmpeg\n"
        "  Windows: winget install ffmpeg"
    )
    sys.exit(1)

if args.audio and not os.path.isfile(args.audio):
    logger.error(f"Audio file not found: {args.audio}")
    sys.exit(1)

openai_api_key = os.environ.get("OPENAI_API_KEY")
if not openai_api_key:
    logger.error("OPENAI_API_KEY environment variable is required")
    sys.exit(1)

openai_client = AsyncOpenAI(api_key=openai_api_key)

youtube_api_key = os.environ.get("YOUTUBE_API_KEY")
if args.youtube_video_id and not youtube_api_key:
    logger.error(
        "YOUTUBE_API_KEY environment variable is required when "
        "--youtube-video-id is set"
    )
    sys.exit(1)

# =============================================================================
# Load topics
# =============================================================================


def load_topics(path: str) -> list[str]:
    """Load story topics from a text file (one per line, # comments)."""
    p = Path(path)
    if not p.is_file():
        logger.error(f"Topics file not found: {path}")
        sys.exit(1)
    topics = [
        line.strip()
        for line in p.read_text(encoding="utf-8").splitlines()
        if line.strip() and not line.strip().startswith("#")
    ]
    if not topics:
        logger.error(f"No topics found in: {path}")
        sys.exit(1)
    return topics


topics = load_topics(args.topics_file)
random.shuffle(topics)
logger.info(f"Loaded {len(topics)} story topics")

# =============================================================================
# Global state (persists across reconnections)
# =============================================================================

# ffmpeg — started once on the first frame; kept alive across sessions so the
# YouTube RTMP stream is uninterrupted even when Reactor reconnects.
ffmpeg_process: subprocess.Popen[bytes] | None = None
frames_sent: int = 0
streaming: bool = False
should_exit: bool = False

# Pipeline bookkeeping
topic_index: int = 0
stories_completed: int = 0
# Pre-generated prompts ready for the next story (set by the pipeline).
next_story_prompts: list[dict] | None = None

# =============================================================================
# YouTube live chat reader
# =============================================================================

YOUTUBE_API_BASE = "https://www.googleapis.com/youtube/v3"


class YouTubeChatReader:
    """
    Reads live chat messages from a YouTube stream via the Data API v3.

    Call :meth:`initialize` once to resolve the ``liveChatId``, then call
    :meth:`poll` periodically (or before each topic pick) to fetch new
    messages.  :meth:`get_fresh_topic` returns a combined topic string from
    the latest *N* messages — or ``None`` if those messages were already used.
    """

    def __init__(self, video_id: str, api_key: str) -> None:
        self.video_id = video_id
        self.api_key = api_key
        self.live_chat_id: str | None = None
        self._next_page_token: str | None = None
        self._polling_interval_s: float = 6.0
        self._messages: list[dict] = []  # [{id, text, author}, ...]
        self._consumed_up_to: int = 0  # index of first unconsumed message
        self._session: aiohttp.ClientSession | None = None

    async def initialize(self) -> bool:
        """
        Resolve the ``liveChatId`` for the configured video.

        Returns ``True`` on success, ``False`` if the video has no active
        live chat (not a livestream, or stream ended).
        """
        self._session = aiohttp.ClientSession()
        try:
            url = (
                f"{YOUTUBE_API_BASE}/videos"
                f"?part=liveStreamingDetails&id={self.video_id}"
                f"&key={self.api_key}"
            )
            async with self._session.get(url) as resp:
                if resp.status != 200:
                    body = await resp.text()
                    logger.error(
                        f"[YTChat] Failed to fetch video info "
                        f"(HTTP {resp.status}): {body[:300]}"
                    )
                    return False
                data = await resp.json()

            items = data.get("items", [])
            if not items:
                logger.error(
                    f"[YTChat] Video not found: {self.video_id}"
                )
                return False

            live_details = items[0].get("liveStreamingDetails", {})
            self.live_chat_id = live_details.get("activeLiveChatId")
            if not self.live_chat_id:
                logger.warning(
                    "[YTChat] No active live chat found for this video. "
                    "Is it a current livestream?"
                )
                return False

            logger.info(
                f"[YTChat] Connected to live chat "
                f"(chatId: {self.live_chat_id[:20]}...)"
            )
            return True

        except Exception as e:
            logger.error(f"[YTChat] Initialization failed: {e}")
            return False

    async def poll(self) -> None:
        """Fetch the latest batch of chat messages."""
        if not self.live_chat_id or not self._session:
            return

        try:
            url = (
                f"{YOUTUBE_API_BASE}/liveChat/messages"
                f"?liveChatId={self.live_chat_id}"
                f"&part=snippet,authorDetails"
                f"&key={self.api_key}"
            )
            if self._next_page_token:
                url += f"&pageToken={self._next_page_token}"

            async with self._session.get(url) as resp:
                if resp.status != 200:
                    body = await resp.text()
                    logger.warning(
                        f"[YTChat] Poll failed (HTTP {resp.status}): "
                        f"{body[:200]}"
                    )
                    return
                data = await resp.json()

            self._next_page_token = data.get("nextPageToken")
            interval_ms = data.get("pollingIntervalMillis", 6000)
            self._polling_interval_s = max(interval_ms / 1000.0, 2.0)

            for item in data.get("items", []):
                snippet = item.get("snippet", {})
                author = item.get("authorDetails", {})
                msg_id = item.get("id", "")

                # Skip if we already have this message
                if any(m["id"] == msg_id for m in self._messages[-50:]):
                    continue

                text = snippet.get("displayMessage") or snippet.get(
                    "textMessageDetails", {}
                ).get("messageText", "")
                if not text.strip():
                    continue

                # Only accept messages that start with the !p command
                stripped = text.strip()
                if not stripped.startswith("!p "):
                    continue
                prompt_text = stripped[3:].strip()
                if not prompt_text:
                    continue

                self._messages.append(
                    {
                        "id": msg_id,
                        "text": prompt_text,
                        "author": author.get("displayName", "?"),
                    }
                )

            # Keep only the last 200 messages in memory
            if len(self._messages) > 200:
                trim = len(self._messages) - 200
                self._messages = self._messages[-200:]
                self._consumed_up_to = max(0, self._consumed_up_to - trim)

        except Exception as e:
            logger.warning(f"[YTChat] Poll error: {e}")

    def get_fresh_topic(
        self, minimum: int = 3, maximum: int = 7
    ) -> str | None:
        """
        Build a topic string from recent unused chat messages.

        Requires at least *minimum* new (not-yet-consumed) messages.
        Consumes up to *maximum* messages at a time.  Returns ``None``
        when there aren't enough new messages since the last call.
        """
        # How many messages are new (arrived after the last consumed index)?
        new_count = len(self._messages) - self._consumed_up_to
        if new_count < minimum:
            return None

        # Take up to `maximum` of the newest messages, but only the new ones
        take = min(new_count, maximum)
        batch = self._messages[self._consumed_up_to : self._consumed_up_to + take]
        self._consumed_up_to += take

        lines = [f'{m["author"]}: {m["text"]}' for m in batch]
        topic = (
            "The audience is chatting live.  Base the visual story on what "
            "they are saying.  Here are the latest messages:\n"
            + "\n".join(lines)
        )
        logger.info(
            f"[YTChat] Using {len(batch)} new chat messages as topic:\n"
            + "\n".join(f"  - {l}" for l in lines)
        )
        return topic

    async def close(self) -> None:
        if self._session:
            await self._session.close()
            self._session = None


# Instantiate the reader (or None if YouTube chat is disabled)
chat_reader: YouTubeChatReader | None = None

if args.youtube_video_id and youtube_api_key:
    chat_reader = YouTubeChatReader(args.youtube_video_id, youtube_api_key)


# =============================================================================
# Topic picker (chat-first, fallback to topics.txt)
# =============================================================================


async def pick_next_topic() -> str:
    """
    Return the next story topic.

    Checks YouTube live chat first (if configured).  If the last 5 messages
    haven't been used yet, they become the topic.  Otherwise falls back to
    the rotating ``topics.txt`` list.
    """
    global topic_index

    if chat_reader:
        try:
            await chat_reader.poll()
            chat_topic = chat_reader.get_fresh_topic(
                minimum=args.chat_messages, maximum=7
            )
            if chat_topic:
                return chat_topic
            logger.debug(
                "[YTChat] No fresh chat messages — falling back to topics.txt"
            )
        except Exception as e:
            logger.warning(f"[YTChat] Error reading chat: {e}")

    topic = topics[topic_index % len(topics)]
    topic_index += 1
    return topic


# =============================================================================
# Story generation (OpenAI)
# =============================================================================

STORY_SYSTEM_PROMPT = """\
You are a cinematic visual story designer.  Given a topic, create a short
visual story told in exactly 7 scenes.  Each scene must be a vivid, detailed
visual description suitable for AI video generation.  The scenes should flow
naturally as a story with clear visual transitions between them.

CRITICAL — Self-contained scene descriptions:
The video model uses ONLY the current scene description for visual conditioning.
It has NO memory of previous scenes.  Therefore every scene description MUST be
fully self-contained: it must re-describe the entire visual setting, environment,
subjects, lighting, and mood — even if they haven't changed.  If a scene
introduces a new element (e.g. a character arrives), the description must ALSO
repeat the existing environment and subjects so they are preserved.

Think of each description as a standalone photograph caption that must reproduce
the full picture.  Omitting any element means it will vanish from the video.

Return ONLY valid JSON in this exact format:
{
    "title": "Story title",
    "scenes": [
        {"scene": 1, "description": "Full, self-contained visual description..."},
        {"scene": 2, "description": "Full, self-contained visual description..."},
        {"scene": 3, "description": "Full, self-contained visual description..."},
        {"scene": 4, "description": "Full, self-contained visual description..."},
        {"scene": 5, "description": "Full, self-contained visual description..."},
        {"scene": 6, "description": "Full, self-contained visual description..."},
        {"scene": 7, "description": "Full, self-contained visual description..."}
    ]
}

Guidelines for scene descriptions:
- Each description should be 2-4 sentences and FULLY self-contained.
- ALWAYS re-describe the environment, setting, subjects, and atmosphere — even
  if they carry over from the previous scene.  Redundancy is required.
- When a new element is introduced, keep describing everything that was already
  in the scene so nothing disappears.
- Describe camera angle, lighting, mood, colors, and key visual elements.
- Make transitions between scenes feel like a natural story progression.
- Include style cues: "cinematic", "high quality", "detailed", "4K".
- Focus only on what the camera SEES — no text, UI elements, or abstractions.\
"""


async def generate_story(topic: str) -> list[dict]:
    """
    Call OpenAI to produce 7 scene descriptions for *topic*.

    Returns a list of dicts: ``[{"frame": int, "prompt": str}, ...]``
    sorted by frame number.
    """
    logger.info(f"[OpenAI] Generating story for: '{topic}'")

    try:
        response = await openai_client.chat.completions.create(
            model=args.openai_model,
            messages=[
                {"role": "system", "content": STORY_SYSTEM_PROMPT},
                {
                    "role": "user",
                    "content": f"Create a 7-scene visual story about: {topic}",
                },
            ],
            temperature=0.9,
            max_tokens=2000,
            response_format={"type": "json_object"},
        )

        content = response.choices[0].message.content
        data = json.loads(content)

        title = data.get("title", topic)
        scenes = data.get("scenes", [])

        # Pad / trim to exactly 7 scenes
        while len(scenes) < NUM_SCENES:
            scenes.append(
                {
                    "scene": len(scenes) + 1,
                    "description": (
                        "A beautiful cinematic landscape, vivid colors, "
                        "high quality, detailed, 4K"
                    ),
                }
            )
        scenes = scenes[:NUM_SCENES]

        story_prompts = []
        for i, (frame, scene) in enumerate(zip(SCENE_FRAMES, scenes)):
            prompt_text = scene.get(
                "description",
                "Cinematic landscape, high quality, detailed, 4K",
            )
            story_prompts.append({"frame": frame, "prompt": prompt_text})
            logger.info(f"  Scene {i + 1} (frame {frame:>3}): {prompt_text[:70]}...")

        logger.info(f"[OpenAI] Story ready: '{title}' ({len(story_prompts)} scenes)")
        return story_prompts

    except Exception as e:
        logger.error(f"[OpenAI] Story generation failed: {e} — using fallback")
        return [
            {
                "frame": f,
                "prompt": (
                    f"A beautiful cinematic scene, chapter {i + 1}, "
                    "vivid colors, high quality, detailed, 4K"
                ),
            }
            for i, f in enumerate(SCENE_FRAMES)
        ]


# =============================================================================
# ffmpeg RTMP streamer
# =============================================================================


def start_ffmpeg(input_width: int, input_height: int) -> subprocess.Popen[bytes]:
    """Spawn an ffmpeg process that encodes raw RGB24 stdin to RTMP/FLV."""
    logger.info(f"Starting ffmpeg -> {args.rtmp_url}")

    ffmpeg_path = shutil.which("ffmpeg")
    if ffmpeg_path is None:
        raise FileNotFoundError("ffmpeg executable not found in PATH")

    cmd: list[str] = [
        ffmpeg_path, "-y",
        # ---- video input (raw frames from stdin) ----
        "-f", "rawvideo",
        "-vcodec", "rawvideo",
        "-pix_fmt", "rgb24",
        "-s", f"{input_width}x{input_height}",
        "-r", str(args.fps),
        "-i", "-",
    ]

    # ---- audio input ----
    if args.audio:
        escaped = (
            args.audio
            .replace("\\", "\\\\")
            .replace(":", "\\:")
            .replace("'", "\\'")
        )
        cmd.extend([
            "-f", "lavfi",
            "-i", f"amovie={escaped}:loop=0,asetpts=N/SR/TB",
            "-map", "0:v",
            "-map", "1:a",
        ])
    else:
        # Silent audio track (required by YouTube / Twitch)
        cmd.extend([
            "-f", "lavfi",
            "-i", "anullsrc=channel_layout=stereo:sample_rate=44100",
            "-map", "0:v",
            "-map", "1:a",
        ])

    # ---- encoding & output ----
    cmd.extend([
        "-c:v", "libx264",
        "-preset", "veryfast",
        "-tune", "zerolatency",
        "-pix_fmt", "yuv420p",
        "-g", str(args.fps * 2),  # keyframe every ~2 s
        "-b:v", "2500k",
        "-maxrate", "3000k",
        "-bufsize", "5000k",
        "-c:a", "aac",
        "-b:a", "128k",
        "-ar", "44100",
        "-f", "flv",
        args.rtmp_url,
    ])

    return subprocess.Popen(
        cmd,
        stdin=subprocess.PIPE,
        stdout=subprocess.DEVNULL,
        stderr=subprocess.PIPE,
    )


# =============================================================================
# Livecore helpers
# =============================================================================


async def schedule_and_start(reactor: Reactor, prompts: list[dict]) -> None:
    """Schedule all scene prompts on livecore and start generation."""
    logger.info(f"Scheduling {len(prompts)} scene prompts on livecore...")
    for p in prompts:
        await reactor.send_command(
            "schedule_prompt",
            {"new_prompt": p["prompt"], "timestamp": p["frame"]},
        )
        logger.debug(f"  Scheduled frame {p['frame']}: {p['prompt'][:50]}...")

    logger.info("All prompts scheduled — starting generation")
    await reactor.send_command("start", {})


# =============================================================================
# Single Reactor session
# =============================================================================


async def run_session(initial_prompts: list[dict] | None = None) -> None:
    """
    Run one Reactor session.  Generates stories in a loop, pipelining the
    next story while the current one renders.

    Returns normally when the connection drops (the outer loop will reconnect).
    Sets ``should_exit`` on fatal errors (ffmpeg permanently dead, etc.).
    """
    global ffmpeg_process, frames_sent, streaming, should_exit
    global topic_index, stories_completed, next_story_prompts

    # -- per-session state --
    generation_complete = asyncio.Event()
    was_generating = False  # tracks whether we've seen paused=False this cycle
    ffmpeg_restart_failures = 0

    # ---- create a fresh Reactor ----
    reactor = Reactor(
        model_name=args.model,
        api_key=args.api_key,
        local=args.local,
    )

    # ---- event handlers (scoped to this session) ----

    @reactor.on_status
    def on_status(status: ReactorStatus) -> None:
        logger.info(f"[Reactor] Status: {status.value}")

    @reactor.on_message
    def on_message(message: object, scope: MessageScope) -> None:
        nonlocal was_generating

        if not isinstance(message, dict):
            logger.debug(f"[Reactor] Message ({scope}): {message}")
            return

        msg_type = message.get("type")

        # -- state updates --
        if msg_type == "state":
            data = message.get("data", {})
            current_frame: int = data.get("current_frame", -1)
            paused: bool = data.get("paused", True)
            current_prompt = data.get("current_prompt") or ""

            if not paused:
                was_generating = True
                if current_frame % 30 == 0:
                    logger.info(
                        f"  Frame {current_frame}/{TOTAL_FRAMES} | "
                        f"{current_prompt[:45]}..."
                    )
            elif paused and was_generating:
                # Generation cycle just finished — model is waiting again
                logger.info(
                    f"  Generation complete (frame {current_frame})"
                )
                was_generating = False
                generation_complete.set()

        # -- discrete events --
        elif msg_type == "event":
            data = message.get("data", {})
            event_name = data.get("event", "unknown")

            if event_name == "prompt_switched":
                frame = data.get("frame", "?")
                new_prompt = data.get("new_prompt", "?")
                logger.info(
                    f"  >>> Transition at frame {frame}: "
                    f"{new_prompt[:55]}..."
                )
            elif event_name == "generation_started":
                logger.info("  Generation started")
            elif event_name == "error":
                logger.error(f"  Model error: {data.get('message', '?')}")
            else:
                logger.info(f"  Event: {event_name}")

        else:
            logger.debug(f"[Reactor] Message: {message}")

    @reactor.on_error
    def on_error(error: ReactorError) -> None:
        logger.error(f"[Reactor] Error: {error}")

    @reactor.on_stream
    def on_stream(track: object) -> None:
        logger.info(f"[Reactor] Video track received: {track}")

    @reactor.on_frame
    def on_frame(frame: NDArray[np.uint8]) -> None:
        nonlocal ffmpeg_restart_failures
        global ffmpeg_process, frames_sent, streaming, should_exit

        if not streaming or should_exit:
            return

        h, w = frame.shape[:2]

        # Enforce the expected model output resolution.  A dimension
        # mismatch between the numpy array and what ffmpeg expects causes
        # raw-video row-stride errors that appear as "TV static".
        if h != MODEL_HEIGHT or w != MODEL_WIDTH:
            if frames_sent == 0:
                logger.warning(
                    f"[ffmpeg] Frame size {w}x{h} differs from expected "
                    f"{MODEL_WIDTH}x{MODEL_HEIGHT} — resizing"
                )
            from PIL import Image as _PILImage

            frame = np.array(
                _PILImage.fromarray(frame).resize(
                    (MODEL_WIDTH, MODEL_HEIGHT), _PILImage.LANCZOS
                )
            )

        # Guarantee C-contiguous memory layout.  VideoFrame.to_ndarray()
        # can return arrays with row-padding; tobytes() would then include
        # those padding bytes, shifting every scanline and producing
        # garbled output.
        if not frame.flags["C_CONTIGUOUS"]:
            frame = np.ascontiguousarray(frame)

        # Lazily (re-)start ffmpeg on the first frame or after it dies
        if ffmpeg_process is None or ffmpeg_process.poll() is not None:
            if ffmpeg_process is not None:
                # Previous process died — log and clean up
                logger.warning("[ffmpeg] Process died, restarting...")
                if ffmpeg_process.stderr:
                    try:
                        stderr = ffmpeg_process.stderr.read().decode()
                        if stderr:
                            logger.error(f"[ffmpeg] stderr:\n{stderr[-500:]}")
                    except Exception:
                        pass
                ffmpeg_process = None

            try:
                ffmpeg_process = start_ffmpeg(MODEL_WIDTH, MODEL_HEIGHT)
                ffmpeg_restart_failures = 0
                logger.info(f"[ffmpeg] Started ({MODEL_WIDTH}x{MODEL_HEIGHT})")
            except Exception as e:
                ffmpeg_restart_failures += 1
                logger.error(
                    f"[ffmpeg] Failed to start ({ffmpeg_restart_failures}): {e}"
                )
                if ffmpeg_restart_failures >= 5:
                    logger.error("[ffmpeg] Too many restart failures — exiting")
                    should_exit = True
                return

        if ffmpeg_process.stdin is None:
            return

        try:
            ffmpeg_process.stdin.write(frame.tobytes())
            frames_sent += 1
            if frames_sent % 300 == 0:
                logger.info(f"[ffmpeg] Frames streamed: {frames_sent}")
        except (BrokenPipeError, OSError):
            # ffmpeg died mid-write — set to None so next frame restarts it
            ffmpeg_process = None

    # ---- connect ----
    logger.info("[Reactor] Connecting...")
    await reactor.connect()
    logger.info("[Reactor] Connected — starting story pipeline")
    streaming = True

    try:
        # Prepare the first story if we don't already have one
        current_prompts = initial_prompts
        if current_prompts is None:
            first_topic = await pick_next_topic()
            current_prompts = await generate_story(first_topic)

        # --- main story loop ---
        while (
            reactor.get_status() != ReactorStatus.DISCONNECTED
            and not should_exit
        ):
            # 1. Schedule prompts and start this story
            generation_complete.clear()
            was_generating = False
            await schedule_and_start(reactor, current_prompts)

            # 2. Pipeline: pick the next topic (chat-first) and generate
            #    the next story in the background while this one renders.
            async def _pipeline_generate_next() -> list[dict]:
                next_topic = await pick_next_topic()
                logger.info(
                    f"[Pipeline] Preparing next story: "
                    f"'{next_topic[:80]}...'"
                )
                return await generate_story(next_topic)

            next_task = asyncio.create_task(_pipeline_generate_next())

            # 3. Wait for generation to finish (or disconnection)
            while (
                not generation_complete.is_set()
                and not should_exit
                and reactor.get_status() != ReactorStatus.DISCONNECTED
            ):
                await asyncio.sleep(0.5)

            if should_exit or reactor.get_status() == ReactorStatus.DISCONNECTED:
                next_task.cancel()
                try:
                    await next_task
                except (asyncio.CancelledError, Exception):
                    pass
                break

            stories_completed += 1
            logger.info(
                f"=== Story #{stories_completed} complete — "
                f"total frames: {frames_sent} ==="
            )

            # 4. Retrieve the pre-generated next story
            try:
                current_prompts = await next_task
            except Exception as e:
                logger.error(f"[Pipeline] Next story failed: {e}. Regenerating...")
                fallback_topic = await pick_next_topic()
                current_prompts = await generate_story(fallback_topic)

            # Brief pause before the next generation starts.  The model has
            # already re-entered its waiting state, so we can schedule
            # immediately.
            await asyncio.sleep(0.5)

    except asyncio.CancelledError:
        pass
    except Exception as e:
        logger.error(f"[Session] Unexpected error: {e}")
    finally:
        streaming = False
        try:
            await reactor.disconnect()
        except Exception:
            pass
        logger.info("[Session] Ended")


# =============================================================================
# Main entry point (reconnection loop)
# =============================================================================


async def main() -> None:
    global should_exit, ffmpeg_process

    logger.info("=" * 62)
    logger.info("  AI Story Livestream")
    logger.info("=" * 62)
    logger.info(f"  Model:        {args.model}")
    logger.info(f"  RTMP URL:     {args.rtmp_url}")
    logger.info(f"  Topics:       {len(topics)} loaded")
    logger.info(f"  Audio:        {args.audio or '(silent)'}")
    logger.info(f"  FPS:          {args.fps}")
    logger.info(f"  OpenAI model: {args.openai_model}")
    logger.info(
        f"  YT Chat:      "
        f"{args.youtube_video_id or 'disabled (no --youtube-video-id)'}"
    )
    logger.info("=" * 62)

    # Initialize YouTube chat reader if configured
    if chat_reader:
        ok = await chat_reader.initialize()
        if ok:
            logger.info("[YTChat] Live chat active — stories will use chat messages")
        else:
            logger.warning(
                "[YTChat] Could not connect to live chat — "
                "falling back to topics.txt only"
            )

    reconnect_count = 0

    while not should_exit and reconnect_count < MAX_RECONNECT_ATTEMPTS:
        try:
            await run_session(initial_prompts=next_story_prompts)
        except Exception as e:
            logger.error(f"[Main] Session crashed: {e}")

        if should_exit:
            break

        reconnect_count += 1
        logger.warning(
            f"[Main] Reconnecting in {RECONNECT_DELAY_S}s "
            f"(attempt {reconnect_count}/{MAX_RECONNECT_ATTEMPTS})..."
        )
        await asyncio.sleep(RECONNECT_DELAY_S)

    # ---- cleanup ----
    if chat_reader:
        await chat_reader.close()

    if ffmpeg_process is not None:
        logger.info("[ffmpeg] Shutting down...")
        if ffmpeg_process.stdin:
            try:
                ffmpeg_process.stdin.close()
            except Exception:
                pass
        ffmpeg_process.terminate()

    logger.info(
        f"Done.  Stories completed: {stories_completed} | "
        f"Frames streamed: {frames_sent}"
    )


if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        logger.info("Interrupted by user")
