#!/usr/bin/env python3
"""
Endless World Explorer — RTMP livestream example for the Reactor Python SDK.

Connects to the hy-world model and autonomously explores the generated world
by sending random movement/camera inputs. Streams the output to an RTMP server
(e.g. YouTube Live, Twitch, or a local RTMP endpoint).

On every generation cycle reset (every 60 chunks), a new random starting image
is picked from the `starting_images/` directory, along with its matching prompt
from a `<name>.prompt.txt` sidecar file.

Requirements:
    - ffmpeg must be installed and available in PATH
    - Place at least one image (.jpg, .png, .webp) in the starting_images/ folder
    - Optionally add a <imagename>.prompt.txt next to each image for per-scene prompts

Usage:
    python main.py --local --rtmp-url rtmp://localhost/live/stream
    python main.py --api-key KEY --model hy-world --rtmp-url rtmp://...
    python main.py --local --rtmp-url rtmp://... --images-dir ./my_images
"""

from __future__ import annotations

import argparse
import asyncio
import base64
import io
import logging
import os
import random
import shutil
import subprocess
import sys
import warnings
from dataclasses import dataclass
from pathlib import Path

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

# =============================================================================
# Configuration
# =============================================================================

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

# Suppress noisy 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)

# Supported image extensions
IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".webp", ".bmp"}

# Movement action choices and their weights (heavily favour forward movement)
KEYBOARD_ACTIONS = ["w", "w", "w", "w", "w", "w", "w", "a", "d", "still"]
CAMERA_ACTIONS = ["still", "still", "still", "left", "right", "up", "down"]

# How often (in seconds) to re-send the current actions to the model.
# Frequent sends ensure we get state messages back for cycle detection.
ACTION_SEND_INTERVAL = 0.5

# How often (in seconds) to pick NEW random actions.
ACTION_CHANGE_INTERVAL_MIN = 2.0
ACTION_CHANGE_INTERVAL_MAX = 6.0

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

parser = argparse.ArgumentParser(
    description="Endless world exploration RTMP livestream"
)
parser.add_argument("--model", "-m", default="hy-world", help="Model name")
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 with stream key"
)
parser.add_argument("--fps", "-f", type=int, default=30, help="Frame rate")
parser.add_argument(
    "--images-dir",
    default=os.path.join(os.path.dirname(__file__), "starting_images"),
    help="Directory with starting images (default: ./starting_images)",
)
parser.add_argument(
    "--prompt",
    default="A beautiful open world landscape, cinematic, high quality, detailed",
    help="Generation prompt",
)
parser.add_argument(
    "--movement-speed",
    type=float,
    default=0.8,
    help="Movement speed 0.0-1.0 (default: 0.8)",
)
parser.add_argument(
    "--rotation-speed",
    type=float,
    default=0.2,
    help="Rotation speed 0.0-1.0 (default: 0.2)",
)
parser.add_argument(
    "--audio", "-a", help="Path to MP3 file to loop as background audio"
)
parser.add_argument(
    "--verbose", "-v", action="store_true", help="Debug logging"
)
args = parser.parse_args()

if args.verbose:
    logger.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"
        "  Windows: winget install ffmpeg\n"
        "  macOS: brew install ffmpeg\n"
        "  Linux: apt install ffmpeg"
    )
    sys.exit(1)

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

# =============================================================================
# Load starting images + per-image prompts
# =============================================================================


@dataclass
class SceneEntry:
    """An image paired with its generation prompt."""
    image_path: Path
    prompt: str


def load_prompt_for_image(image_path: Path, fallback_prompt: str) -> str:
    """
    Load a prompt from a sidecar .prompt.txt file next to the image.

    Looks for `<stem>.prompt.txt` in the same directory as the image.
    Falls back to the CLI --prompt value if no sidecar file exists.
    """
    prompt_file = image_path.parent / f"{image_path.stem}.prompt.txt"
    if prompt_file.is_file():
        text = prompt_file.read_text(encoding="utf-8").strip()
        if text:
            return text
    return fallback_prompt


def load_scenes(images_dir: str, fallback_prompt: str) -> list[SceneEntry]:
    """Discover all image files and their paired prompts."""
    dirpath = Path(images_dir)
    if not dirpath.is_dir():
        logger.warning(f"Images directory not found: {images_dir}")
        return []
    entries: list[SceneEntry] = []
    for p in sorted(dirpath.iterdir()):
        if p.is_file() and p.suffix.lower() in IMAGE_EXTENSIONS:
            prompt = load_prompt_for_image(p, fallback_prompt)
            entries.append(SceneEntry(image_path=p, prompt=prompt))
    return entries


MAX_BASE64_BYTES = 240 * 1024  # 240 KB limit on the base64 string itself


def image_to_base64(image_path: Path) -> str:
    """
    Read an image, compress it so the base64 output stays under MAX_BASE64_BYTES.

    If the raw file is already small enough it is sent as-is. Otherwise the
    image is re-encoded as JPEG with decreasing quality until it fits.
    """
    from PIL import Image

    raw = image_path.read_bytes()

    # Base64 inflates by 4/3, so raw bytes must be under 3/4 of the b64 limit
    raw_limit = int(MAX_BASE64_BYTES * 3 / 4)

    if len(raw) <= raw_limit:
        return base64.b64encode(raw).decode("utf-8")

    # Need to compress — open with PIL and reduce quality
    img = Image.open(image_path)
    img = img.convert("RGB")

    # Optionally downscale very large images first (max 1280 on longest side)
    max_side = 1280
    if max(img.size) > max_side:
        img.thumbnail((max_side, max_side), Image.LANCZOS)

    for quality in (85, 75, 60, 45, 30):
        buf = io.BytesIO()
        img.save(buf, format="JPEG", quality=quality, optimize=True)
        data = buf.getvalue()
        if len(data) <= raw_limit:
            logger.debug(
                f"Compressed {image_path.name}: {len(raw)}B -> {len(data)}B (q={quality})"
            )
            return base64.b64encode(data).decode("utf-8")

    # Last resort: use the lowest quality attempt
    logger.warning(
        f"{image_path.name} still {len(data)}B after compression (limit {raw_limit}B)"
    )
    return base64.b64encode(data).decode("utf-8")


scenes = load_scenes(args.images_dir, args.prompt)
if scenes:
    logger.info(
        f"Found {len(scenes)} starting scene(s) in {args.images_dir}"
    )
    for s in scenes:
        has_custom = s.prompt != args.prompt
        tag = " (custom prompt)" if has_custom else " (default prompt)"
        logger.info(f"  - {s.image_path.name}{tag}")
        if has_custom:
            logger.info(f"    prompt: {s.prompt[:80]}...")
else:
    logger.warning(
        "No starting images found. The model will use its default image. "
        f"Place images in: {args.images_dir}"
    )

# =============================================================================
# Initialize Reactor
# =============================================================================

reactor = Reactor(
    model_name=args.model,
    api_key=args.api_key,
    local=args.local,
)

# =============================================================================
# State
# =============================================================================

frames_sent = 0
streaming = False
ffmpeg_process: subprocess.Popen[bytes] | None = None
should_exit = False

# Track the last known chunk_index so we can detect resets
last_chunk_index: int = -1
current_scene_index: int = 0  # index into scenes list (cycles through)
generation_cycle: int = 0  # how many full 60-chunk cycles we've completed
current_prompt: str = args.prompt  # active prompt (updated per scene)

# =============================================================================
# Event handlers
# =============================================================================


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


@reactor.on_message
def on_message(message: object, scope: MessageScope) -> None:
    """Handle messages from the model, including state updates."""
    global last_chunk_index

    if isinstance(message, dict) and message.get("type") == "state":
        data = message.get("data", {})
        chunk_index = data.get("chunk_index", -1)

        # Detect cycle reset: chunk_index decreased significantly.
        # The model emits state only when actions change, so we might
        # not see the exact 59->0 transition. Any big drop means a reset.
        if last_chunk_index >= 0 and chunk_index < last_chunk_index:
            logger.info(
                f"Chunk index dropped {last_chunk_index} -> {chunk_index}, "
                f"detected generation cycle reset"
            )
            _on_generation_cycle_reset()

        last_chunk_index = chunk_index

        # Log every chunk at debug, every 10th at info
        kb = data.get("current_keyboard_action", "?")
        cam = data.get("current_camera_action", "?")
        if chunk_index % 10 == 0:
            logger.info(f"Chunk {chunk_index}/59 | move={kb} cam={cam}")
        else:
            logger.debug(f"Chunk {chunk_index}/59 | move={kb} cam={cam}")
    else:
        logger.debug(f"Message: {message}")


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


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


@reactor.on_frame
def on_frame(frame: NDArray[np.uint8]) -> None:
    """Send each frame to ffmpeg for RTMP streaming."""
    global frames_sent, ffmpeg_process, streaming, should_exit

    if not streaming or should_exit:
        return

    # Start ffmpeg on first frame (now we know input dimensions)
    if ffmpeg_process is None:
        h, w = frame.shape[:2]
        try:
            ffmpeg_process = start_ffmpeg(w, h)
            logger.info(f"ffmpeg started ({w}x{h})")
        except Exception as e:
            logger.error(f"Failed to start ffmpeg: {e}")
            should_exit = True
            return

    if ffmpeg_process.stdin is None or ffmpeg_process.poll() is not None:
        return

    try:
        ffmpeg_process.stdin.write(frame.tobytes())
        frames_sent += 1
        if frames_sent % 300 == 0:
            logger.info(f"Frames sent: {frames_sent}")
    except (BrokenPipeError, OSError):
        pass  # ffmpeg exited, will be handled in main loop


# =============================================================================
# Exploration logic
# =============================================================================


def _on_generation_cycle_reset() -> None:
    """Called when the model finishes 60 chunks and resets to chunk 0."""
    global generation_cycle, current_scene_index

    generation_cycle += 1
    logger.info(
        f"=== Generation cycle {generation_cycle} complete — starting new cycle ==="
    )

    # Pick the next scene (round-robin with shuffle on wrap)
    if scenes:
        current_scene_index = (current_scene_index + 1) % len(scenes)
        if current_scene_index == 0:
            random.shuffle(scenes)

        scene = scenes[current_scene_index]
        logger.info(f"Switching to scene: {scene.image_path.name}")
        logger.info(f"  prompt: {scene.prompt[:80]}...")
        _send_scene(scene)


def _send_scene(scene: SceneEntry) -> None:
    """Encode and send a starting image + its prompt to the model."""
    global current_prompt
    try:
        # Send the prompt first so it's ready when the image triggers a restart
        current_prompt = scene.prompt
        asyncio.get_event_loop().create_task(
            reactor.send_command("set_prompt", {"prompt": scene.prompt})
        )

        b64 = image_to_base64(scene.image_path)
        asyncio.get_event_loop().create_task(
            reactor.send_command("set_starting_image", {"image_base64": b64})
        )
        logger.info(
            f"Sent scene: {scene.image_path.name} ({len(b64)} chars) "
            f"+ prompt ({len(scene.prompt)} chars)"
        )
    except Exception as e:
        logger.error(f"Failed to send scene: {e}")


async def send_random_actions() -> None:
    """
    Continuously send random keyboard and camera actions to explore the world.

    Runs as a background coroutine. Changes actions at random intervals
    to create organic-looking exploration paths.
    """
    logger.info("Random exploration controller started, waiting for READY status...")

    # Wait until reactor is ready before sending commands.
    # Check for both READY and also that we haven't disconnected.
    while True:
        status = reactor.get_status()
        if status == ReactorStatus.READY:
            break
        if status == ReactorStatus.DISCONNECTED or should_exit:
            logger.warning("Exploration controller exiting (disconnected or should_exit)")
            return
        await asyncio.sleep(0.5)

    logger.info("Reactor is READY — sending initial commands")

    # Set movement and rotation speeds
    try:
        logger.info(f"Setting movement_speed={args.movement_speed}, rotation_speed={args.rotation_speed}")
        await reactor.send_command(
            "set_movement_speed", {"speed": args.movement_speed}
        )
        await reactor.send_command(
            "set_rotation_speed", {"speed": args.rotation_speed}
        )

        # Send the first scene (image + its prompt), or just the default prompt
        if scenes:
            first_scene = scenes[0]
            logger.info(f"Sending initial scene: {first_scene.image_path.name}")
            logger.info(f"  prompt: {first_scene.prompt[:80]}...")
            await reactor.send_command(
                "set_prompt", {"prompt": first_scene.prompt}
            )
            b64 = image_to_base64(first_scene.image_path)
            await reactor.send_command(
                "set_starting_image", {"image_base64": b64}
            )
        else:
            logger.info(f"Setting default prompt: {args.prompt[:60]}...")
            await reactor.send_command("set_prompt", {"prompt": args.prompt})
        logger.info("Initial commands sent successfully")
    except Exception as e:
        logger.error(f"Failed to send initial commands: {e}")
        return

    # Camera balance: track cumulative drift on yaw (left/right) and pitch
    # (up/down). When drift on an axis exceeds the threshold, force a
    # corrective action back toward center instead of picking randomly.
    yaw_drift = 0    # positive = drifted right, negative = drifted left
    pitch_drift = 0  # positive = drifted down, negative = drifted up
    DRIFT_THRESHOLD = 1  # correct after a single directional move

    # Start sending actions
    logger.info("Starting random exploration loop")
    send_count = 0

    # Current actions (re-sent frequently, changed less often)
    keyboard = random.choice(KEYBOARD_ACTIONS)
    camera = "still"

    # Timer for when to pick new random actions
    next_change_at = asyncio.get_event_loop().time() + random.uniform(
        ACTION_CHANGE_INTERVAL_MIN, ACTION_CHANGE_INTERVAL_MAX
    )

    while not should_exit:
        try:
            now = asyncio.get_event_loop().time()

            # Time to pick new actions?
            if now >= next_change_at:
                keyboard = random.choice(KEYBOARD_ACTIONS)

                # Decide camera action: correct if drifted too far, else random
                camera = random.choice(CAMERA_ACTIONS)
                if abs(yaw_drift) >= DRIFT_THRESHOLD:
                    camera = "left" if yaw_drift > 0 else "right"
                elif abs(pitch_drift) >= DRIFT_THRESHOLD:
                    camera = "up" if pitch_drift > 0 else "down"

                # Update drift counters
                if camera == "left":
                    yaw_drift -= 1
                elif camera == "right":
                    yaw_drift += 1
                elif camera == "up":
                    pitch_drift -= 1
                elif camera == "down":
                    pitch_drift += 1

                next_change_at = now + random.uniform(
                    ACTION_CHANGE_INTERVAL_MIN, ACTION_CHANGE_INTERVAL_MAX
                )
                send_count += 1
                if send_count <= 3 or send_count % 10 == 0:
                    logger.info(
                        f"New action #{send_count} -> move={keyboard} cam={camera} "
                        f"(yaw={yaw_drift}, pitch={pitch_drift})"
                    )

            # Always re-send current actions (keeps state messages flowing)
            await reactor.send_command(
                "set_keyboard_action", {"action": keyboard}
            )
            await reactor.send_command(
                "set_camera_action", {"action": camera}
            )

            await asyncio.sleep(ACTION_SEND_INTERVAL)

        except Exception as e:
            if not should_exit:
                logger.error(f"Error sending action: {e}")
            await asyncio.sleep(1)


# =============================================================================
# ffmpeg
# =============================================================================


def start_ffmpeg(input_width: int, input_height: int) -> subprocess.Popen[bytes]:
    """Start ffmpeg process for RTMP streaming."""
    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 = [
        ffmpeg_path, "-y",
        # Video input from stdin
        "-f", "rawvideo",
        "-vcodec", "rawvideo",
        "-pix_fmt", "rgb24",
        "-s", f"{input_width}x{input_height}",
        "-r", str(args.fps),
        "-i", "-",
    ]

    # Audio input: either loop an MP3 file or generate silent audio
    if args.audio:
        escaped_path = args.audio.replace("\\", "\\\\").replace(":", "\\:").replace("'", "\\'")
        cmd.extend([
            "-f", "lavfi",
            "-i", f"amovie={escaped_path}:loop=0,asetpts=N/SR/TB",
            "-map", "0:v",
            "-map", "1:a",
        ])
    else:
        # Generate silent audio (required by YouTube/Twitch)
        cmd.extend([
            "-f", "lavfi",
            "-i", "anullsrc=channel_layout=stereo:sample_rate=44100",
            "-map", "0:v",
            "-map", "1:a",
        ])

    cmd.extend([
        # Video encoding
        "-c:v", "libx264",
        "-preset", "veryfast",
        "-tune", "zerolatency",
        "-pix_fmt", "yuv420p",
        "-g", "60",  # Keyframe every 2 seconds at 30fps
        "-b:v", "2500k",
        "-maxrate", "3000k",
        "-bufsize", "5000k",
        # Audio encoding
        "-c:a", "aac",
        "-b:a", "128k",
        "-ar", "44100",
        # Output
        "-f", "flv",
        args.rtmp_url,
    ])

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


# =============================================================================
# Main
# =============================================================================


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

    logger.info("=" * 60)
    logger.info("  Endless World Explorer")
    logger.info("=" * 60)
    logger.info(f"Model:          {args.model}")
    logger.info(f"RTMP URL:       {args.rtmp_url}")
    logger.info(f"Default prompt: {args.prompt[:60]}...")
    logger.info(f"Scenes:         {len(scenes)} found")
    logger.info(f"Audio:          {args.audio or '(silent)'}")
    logger.info(f"Move speed:     {args.movement_speed}")
    logger.info(f"Rotation speed: {args.rotation_speed}")
    logger.info("=" * 60)

    logger.info("Connecting to Reactor...")
    await reactor.connect()
    logger.info("Connected! Waiting for first frame to start ffmpeg...")
    streaming = True

    # Start the random action controller as a background task
    action_task = asyncio.create_task(send_random_actions())

    try:
        while (
            reactor.get_status() != ReactorStatus.DISCONNECTED
            and not should_exit
        ):
            # Check if ffmpeg died
            if ffmpeg_process is not None and ffmpeg_process.poll() is not None:
                if ffmpeg_process.stderr:
                    stderr = ffmpeg_process.stderr.read().decode()
                    if stderr:
                        logger.error(f"ffmpeg error:\n{stderr[-1000:]}")
                break
            # Check if the action controller died unexpectedly
            if action_task.done() and not should_exit:
                exc = action_task.exception() if not action_task.cancelled() else None
                if exc:
                    logger.error(f"Action controller crashed: {exc}")
                else:
                    logger.warning("Action controller exited early, restarting...")
                action_task = asyncio.create_task(send_random_actions())
            await asyncio.sleep(1)
    except asyncio.CancelledError:
        pass

    # Shutdown
    should_exit = True
    streaming = False
    logger.info("Stopping...")

    # Cancel the action controller
    action_task.cancel()
    try:
        await action_task
    except asyncio.CancelledError:
        pass

    # Stop ffmpeg
    if ffmpeg_process is not None:
        if ffmpeg_process.stdin:
            ffmpeg_process.stdin.close()
        ffmpeg_process.terminate()

    await reactor.disconnect()
    logger.info(
        f"Done. Total frames sent: {frames_sent} | "
        f"Generation cycles: {generation_cycle}"
    )


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