#!/usr/bin/env python
"""
RTMP Demo Script

This script demonstrates the RTMP streaming capabilities of FMUS-VID
using a local video file as a source (as a stand-in for a live RTMP stream).

It simulates the streaming process and shows how real-time transcription,
translation, and subtitle overlay would work in a production environment.
"""

import os
import sys
import asyncio
import time
import logging
from pathlib import Path
import numpy as np
import cv2

# Add parent directory to path for imports
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

# Setup logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# Check if required modules are available
try:
    from fmusvid.stream import RTMPCapture, RTMPOutput
    from fmusvid.ai.speech import WhisperTranscriber
    from fmusvid.ai.translation import LLMTranslator
    from fmusvid.operations.text import SubtitleOverlay

    MODULES_AVAILABLE = True
except ImportError as e:
    logger.warning(f"RTMP modules not fully available: {e}")
    MODULES_AVAILABLE = False

# Default values
DEFAULT_VIDEO_PATH = os.path.join("examples", "assets", "sample.mp4")
DEFAULT_OUTPUT_DIR = os.path.join("examples", "output")
DEFAULT_LANGUAGES = ["en", "es", "fr", "de"]
DEFAULT_FPS = 30.0

class VideoFileSimulator:
    """
    Simulates an RTMP stream using a local video file.
    This is a stand-in for the RTMPCapture class when no actual RTMP server is available.
    """
    def __init__(self, video_path, fps=30.0):
        self.video_path = video_path
        self.target_fps = fps
        self.cap = None
        self.is_running = False
        self.frame_count = 0
        self.audio_data = np.zeros((1024, 1), dtype=np.float32)  # Dummy audio

    async def start(self):
        """Start the video capture."""
        self.cap = cv2.VideoCapture(self.video_path)
        if not self.cap.isOpened():
            raise ValueError(f"Could not open video file: {self.video_path}")
        self.is_running = True
        self.frame_count = 0
        logger.info(f"Started video file simulator: {self.video_path}")
        return True

    async def stop(self):
        """Stop the video capture."""
        if self.cap:
            self.cap.release()
        self.is_running = False
        logger.info("Stopped video file simulator")

    async def get_frame(self):
        """Get the next frame from the video."""
        if not self.is_running or not self.cap:
            return None

        ret, frame = self.cap.read()
        if not ret:
            return None

        self.frame_count += 1
        return frame

    async def get_audio(self):
        """Simulate getting audio from the video."""
        # In a real implementation, we would extract audio from the video
        # Here we just return dummy audio data
        return self.audio_data

    async def aiter_frames_with_audio(self):
        """Async iterator for frames and audio."""
        while self.is_running:
            frame = await self.get_frame()
            if frame is None:
                break

            audio = await self.get_audio()

            # Simulate real-time playback
            await asyncio.sleep(1/self.target_fps)

            yield frame, audio


class VideoFileOutput:
    """
    Simulates an RTMP output using a local video file.
    This is a stand-in for the RTMPOutput class when no actual RTMP server is available.
    """
    def __init__(self, output_path, width=1280, height=720, fps=30.0):
        self.output_path = output_path
        self.width = width
        self.height = height
        self.fps = fps
        self.writer = None
        self.is_running = False

    async def start(self):
        """Start the video writer."""
        # Ensure output directory exists
        os.makedirs(os.path.dirname(self.output_path), exist_ok=True)

        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        self.writer = cv2.VideoWriter(
            self.output_path,
            fourcc,
            self.fps,
            (self.width, self.height)
        )
        self.is_running = True
        logger.info(f"Started video file output: {self.output_path}")
        return True

    async def stop(self):
        """Stop the video writer."""
        if self.writer:
            self.writer.release()
        self.is_running = False
        logger.info(f"Stopped video file output: {self.output_path}")

    async def write_frame(self, frame):
        """Write a frame to the output video."""
        if not self.is_running or not self.writer:
            return False

        # Resize frame if needed
        if frame.shape[1] != self.width or frame.shape[0] != self.height:
            frame = cv2.resize(frame, (self.width, self.height))

        self.writer.write(frame)
        return True


class DummyTranscriber:
    """
    Dummy transcriber class for demonstration purposes.
    """
    def __init__(self):
        self.counter = 0
        self.transcripts = [
            "Hello, welcome to the RTMP streaming demo.",
            "This demonstrates real-time transcription and translation.",
            "We're simulating audio transcription right now.",
            "In a real implementation, we would use Whisper for this.",
            "The transcription would be more accurate and responsive.",
            "This is just a placeholder for demonstration purposes."
        ]

    async def transcribe_chunk(self, audio):
        """Simulate transcribing audio chunks."""
        # Only return text every 30 frames (about 1 second at 30fps)
        if self.counter % 30 == 0:
            text = self.transcripts[(self.counter // 30) % len(self.transcripts)]
            self.counter += 1
            return text
        self.counter += 1
        return ""


class DummyTranslator:
    """
    Dummy translator class for demonstration purposes.
    """
    def __init__(self, target_languages=None):
        self.target_languages = target_languages or ["es", "fr", "de"]
        self.translations = {
            "es": {
                "Hello, welcome to the RTMP streaming demo.": "Hola, bienvenido a la demostración de streaming RTMP.",
                "This demonstrates real-time transcription and translation.": "Esto demuestra la transcripción y traducción en tiempo real.",
                "We're simulating audio transcription right now.": "Estamos simulando la transcripción de audio ahora mismo.",
                "In a real implementation, we would use Whisper for this.": "En una implementación real, usaríamos Whisper para esto.",
                "The transcription would be more accurate and responsive.": "La transcripción sería más precisa y receptiva.",
                "This is just a placeholder for demonstration purposes.": "Esto es solo un marcador de posición con fines de demostración."
            },
            "fr": {
                "Hello, welcome to the RTMP streaming demo.": "Bonjour, bienvenue à la démo de streaming RTMP.",
                "This demonstrates real-time transcription and translation.": "Ceci démontre la transcription et la traduction en temps réel.",
                "We're simulating audio transcription right now.": "Nous simulons la transcription audio en ce moment.",
                "In a real implementation, we would use Whisper for this.": "Dans une implémentation réelle, nous utiliserions Whisper pour cela.",
                "The transcription would be more accurate and responsive.": "La transcription serait plus précise et réactive.",
                "This is just a placeholder for demonstration purposes.": "Ce n'est qu'un espace réservé à des fins de démonstration."
            },
            "de": {
                "Hello, welcome to the RTMP streaming demo.": "Hallo, willkommen zur RTMP-Streaming-Demo.",
                "This demonstrates real-time transcription and translation.": "Dies demonstriert Echtzeit-Transkription und -Übersetzung.",
                "We're simulating audio transcription right now.": "Wir simulieren gerade die Audiotranskription.",
                "In a real implementation, we would use Whisper for this.": "In einer realen Implementierung würden wir dafür Whisper verwenden.",
                "The transcription would be more accurate and responsive.": "Die Transkription wäre genauer und reaktionsschneller.",
                "This is just a placeholder for demonstration purposes.": "Dies ist nur ein Platzhalter für Demonstrationszwecke."
            }
        }

    async def translate_async(self, text):
        """Simulate translating text."""
        result = {}
        for lang in self.target_languages:
            if lang in self.translations and text in self.translations[lang]:
                result[lang] = self.translations[lang][text]
            else:
                # Fallback translation
                result[lang] = f"[{lang}] {text}"
        return result


class SimpleSubtitleOverlay:
    """
    Simple subtitle overlay class for demonstration purposes.
    """
    def __init__(self, position="bottom", font_size=28, background_opacity=0.7):
        self.position = position
        self.font_size = font_size
        self.background_opacity = background_opacity

    def apply(self, frame, text):
        """Apply subtitle overlay to frame."""
        if not text:
            return frame.copy()

        # Copy frame to avoid modifying original
        result = frame.copy()

        # Get frame dimensions
        height, width = result.shape[:2]

        # Calculate text size
        font = cv2.FONT_HERSHEY_SIMPLEX
        font_scale = self.font_size / 30.0
        thickness = 1
        (text_width, text_height), baseline = cv2.getTextSize(
            text, font, font_scale, thickness
        )

        # Determine text position
        if self.position == "bottom":
            text_x = (width - text_width) // 2
            text_y = height - 30
        else:
            text_x = (width - text_width) // 2
            text_y = 50

        # Draw background
        pad = 10
        bg_rect = (
            text_x - pad,
            text_y - text_height - pad,
            text_width + 2 * pad,
            text_height + 2 * pad
        )

        overlay = result.copy()
        cv2.rectangle(
            overlay,
            (bg_rect[0], bg_rect[1]),
            (bg_rect[0] + bg_rect[2], bg_rect[1] + bg_rect[3]),
            (0, 0, 0),
            -1
        )

        # Apply overlay with transparency
        cv2.addWeighted(
            overlay, self.background_opacity,
            result, 1 - self.background_opacity,
            0, result
        )

        # Draw text
        cv2.putText(
            result,
            text,
            (text_x, text_y),
            font,
            font_scale,
            (255, 255, 255),
            thickness
        )

        return result


async def process_video_file(
    video_path,
    output_dir,
    languages=None,
    fps=30.0,
    max_duration=60.0  # 1 minute max by default
):
    """
    Process a video file and generate outputs with subtitles in multiple languages.

    Args:
        video_path: Path to the input video file
        output_dir: Directory to save output videos
        languages: List of language codes to generate subtitles for
        fps: Output frame rate
        max_duration: Maximum duration to process in seconds
    """
    languages = languages or DEFAULT_LANGUAGES
    os.makedirs(output_dir, exist_ok=True)

    logger.info(f"Processing video file: {video_path}")
    logger.info(f"Output directory: {output_dir}")
    logger.info(f"Languages: {', '.join(languages)}")

    try:
        # Create video file simulator (instead of RTMP capture)
        video_simulator = VideoFileSimulator(video_path, fps=fps)
        await video_simulator.start()

        # Get video dimensions
        cap = cv2.VideoCapture(video_path)
        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        cap.release()

        # Create dummy transcriber and translator
        transcriber = DummyTranscriber()
        translator = DummyTranslator(target_languages=[lang for lang in languages if lang != "en"])

        # Create subtitle overlays
        subtitle_overlays = {
            lang: SimpleSubtitleOverlay(
                position="bottom",
                font_size=28,
                background_opacity=0.7
            )
            for lang in languages
        }

        # Create video outputs
        video_outputs = {}
        for lang in languages:
            output_path = os.path.join(output_dir, f"output_{lang}.mp4")
            video_outputs[lang] = VideoFileOutput(
                output_path=output_path,
                width=width,
                height=height,
                fps=fps
            )
            await video_outputs[lang].start()

        # Process frames
        start_time = time.time()
        frame_count = 0
        current_text = {"en": ""}
        subtitle_display_time = 0

        # Main processing loop
        async for frame, audio in video_simulator.aiter_frames_with_audio():
            try:
                # Check if we've reached the maximum duration
                elapsed = time.time() - start_time
                if elapsed > max_duration:
                    logger.info(f"Reached maximum duration of {max_duration}s")
                    break

                # Process audio for transcription
                text = await transcriber.transcribe_chunk(audio)
                if text:
                    logger.info(f"Transcribed: {text}")
                    current_text["en"] = text
                    subtitle_display_time = time.time() + 3.0  # Display for 3 seconds

                    # Translate to other languages
                    if text.strip():
                        translations = await translator.translate_async(text)
                        # Update current text for each language
                        for lang, translated_text in translations.items():
                            if lang in languages and lang != "en":
                                current_text[lang] = translated_text
                                logger.info(f"Translated ({lang}): {translated_text}")

                # Process frames for each language
                for lang in languages:
                    # Get text for this language
                    text = current_text.get(lang, "")

                    # Clear subtitle if display time has elapsed
                    if time.time() > subtitle_display_time:
                        text = ""

                    # Apply subtitle overlay
                    if text and lang in subtitle_overlays:
                        frame_with_subtitle = subtitle_overlays[lang].apply(frame, text)
                    else:
                        frame_with_subtitle = frame.copy()

                    # Write frame to output
                    if lang in video_outputs:
                        await video_outputs[lang].write_frame(frame_with_subtitle)

                # Update stats
                frame_count += 1
                if frame_count % 100 == 0:
                    fps_actual = frame_count / (time.time() - start_time)
                    logger.info(f"Processed {frame_count} frames ({fps_actual:.2f} fps)")

            except Exception as e:
                logger.error(f"Error processing frame: {e}")

        # Clean up
        for output in video_outputs.values():
            await output.stop()

        await video_simulator.stop()

        logger.info(f"Processing complete. Processed {frame_count} frames in {time.time() - start_time:.2f}s")
        logger.info(f"Output videos saved to {output_dir}")

        # Print output video paths
        for lang in languages:
            output_path = os.path.join(output_dir, f"output_{lang}.mp4")
            logger.info(f"{lang} output: {output_path}")

    except Exception as e:
        logger.error(f"Error in video processing: {e}")


def main():
    """Main function to run the demo."""
    import argparse

    parser = argparse.ArgumentParser(description="RTMP Demo using local video files")
    parser.add_argument("--video", default=DEFAULT_VIDEO_PATH, help="Path to input video file")
    parser.add_argument("--output", default=DEFAULT_OUTPUT_DIR, help="Output directory")
    parser.add_argument("--languages", default=",".join(DEFAULT_LANGUAGES), help="Comma-separated list of language codes")
    parser.add_argument("--fps", type=float, default=DEFAULT_FPS, help="Output frame rate")
    parser.add_argument("--duration", type=float, default=60.0, help="Maximum duration to process in seconds")

    args = parser.parse_args()

    print("\nRTMP Demo Script\n")
    print("This script demonstrates the RTMP streaming capabilities of FMUS-VID")
    print("using a local video file as a source.")
    print("\nFeatures demonstrated:")
    print("- Video stream processing")
    print("- Simulated real-time audio transcription")
    print("- Simulated translation to multiple languages")
    print("- Subtitle overlay")
    print("- Multi-language video output")

    # Create output directory
    os.makedirs(args.output, exist_ok=True)

    # Check if video file exists
    if not os.path.exists(args.video):
        print(f"\nError: Video file not found: {args.video}")
        print("You can specify a different video file with --video")
        return

    try:
        # Run the async processing function
        asyncio.run(process_video_file(
            video_path=args.video,
            output_dir=args.output,
            languages=args.languages.split(","),
            fps=args.fps,
            max_duration=args.duration
        ))

        print("\nDemo completed successfully!")
        print(f"Output videos saved to: {args.output}")

    except KeyboardInterrupt:
        print("\nDemo stopped by user")
    except Exception as e:
        print(f"\nError running demo: {e}")


if __name__ == "__main__":
    main()
