#!/usr/bin/env python
"""
Real-time RTMP Streaming with Multilingual Subtitles Example for FMUS-VID.

This example demonstrates how to:
1. Capture an RTMP stream
2. Extract and transcribe audio in real-time
3. Translate subtitles into multiple languages using a local LLM
4. Overlay translated subtitles onto the video
5. Output multiple RTMP streams with different language subtitles

Requirements:
- faster-whisper
- ollama (or another local LLM)
- nginx with RTMP module
"""

import os
import sys
import asyncio
import time
from pathlib import Path
import logging
import numpy as np
import cv2
from dataclasses import dataclass
from typing import List, Dict, Optional, Tuple, AsyncGenerator

# Add parent directory to path to allow running from examples directory
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

# This is a conceptual example - these modules would need to be implemented
from fmusvid.stream import RTMPCapture, RTMPOutput, StreamMonitor
from fmusvid.ai.speech import WhisperTranscriber
from fmusvid.ai.translation import LLMTranslator
from fmusvid.operations.text import SubtitleOverlay

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

@dataclass
class SubtitleEntry:
    """A single subtitle entry with timing information."""
    text: str
    start_time: float
    end_time: float
    language: str = "en"

class RTMPSubtitleProcessor:
    """Process RTMP streams and add multilingual subtitles."""

    def __init__(
        self,
        input_rtmp: str,
        output_base_rtmp: str,
        languages: List[str] = ["en", "es", "fr", "de"],
        whisper_model: str = "medium",
        llm_model: str = "ollama:mistral",
        device: str = "cuda",
        subtitle_history_length: int = 10
    ):
        """
        Initialize the RTMP subtitle processor.

        Args:
            input_rtmp: RTMP input URL (e.g., "rtmp://localhost/live/stream")
            output_base_rtmp: Base RTMP output URL (e.g., "rtmp://localhost/live/stream_out")
            languages: List of language codes to generate subtitles for
            whisper_model: Model size for Whisper transcription
            llm_model: LLM model to use for translation
            device: Device to run models on ("cuda" or "cpu")
            subtitle_history_length: Number of recent subtitles to keep in history
        """
        self.input_rtmp = input_rtmp
        self.output_base_rtmp = output_base_rtmp
        self.languages = languages
        self.whisper_model = whisper_model
        self.llm_model = llm_model
        self.device = device
        self.subtitle_history_length = subtitle_history_length

        # Subtitle history for each language
        self.subtitle_history = {lang: [] for lang in languages}

        # Metrics for monitoring
        self.metrics = {
            "transcription_latency": 0,
            "translation_latency": 0,
            "end_to_end_latency": 0,
            "fps": 0
        }

        # Initialize components
        logger.info("Initializing RTMP subtitle processor...")

        # These would be actual component initializations in a real implementation
        self.stream = RTMPCapture(input_rtmp)
        self.transcriber = WhisperTranscriber(model=whisper_model, device=device)
        self.translator = LLMTranslator(
            model=llm_model,
            target_languages=[lang for lang in languages if lang != "en"]
        )
        self.outputs = {
            lang: RTMPOutput(f"{output_base_rtmp}_{lang}")
            for lang in languages
        }
        self.monitor = StreamMonitor(metrics=["bitrate", "fps", "latency"])
        self.monitor.attach_stream(self.stream)

        # Subtitle overlay for each language
        self.subtitle_overlays = {
            lang: SubtitleOverlay(
                position="bottom",
                font_size=24,
                background_opacity=0.7,
                font_color=(255, 255, 255)
            )
            for lang in languages
        }

        logger.info(f"Initialized processor for languages: {', '.join(languages)}")

    async def process(self):
        """Process the RTMP stream and generate subtitled outputs."""
        logger.info(f"Starting to process RTMP stream: {self.input_rtmp}")
        frame_count = 0
        start_time = time.time()

        # In a real implementation, this would use the actual RTMPCapture.aiter_frames_with_audio
        async for frame, audio in self._mock_rtmp_stream():
            try:
                # Measure processing start time
                frame_start_time = time.time()

                # Process audio for transcription
                transcription_start = time.time()
                text = await self.transcriber.transcribe_chunk(audio)
                self.metrics["transcription_latency"] = time.time() - transcription_start

                if text:
                    logger.info(f"Transcribed: {text}")

                    # Create subtitle entry
                    current_time = time.time() - start_time
                    subtitle = SubtitleEntry(
                        text=text,
                        start_time=current_time,
                        end_time=current_time + 3.0,  # Display for 3 seconds
                        language="en"
                    )

                    # Add to English subtitle history
                    self._add_to_subtitle_history("en", subtitle)

                    # Translate to other languages
                    if text.strip():
                        translation_start = time.time()
                        translations = await self.translator.translate_async(text)
                        self.metrics["translation_latency"] = time.time() - translation_start

                        # Add translated subtitles to history
                        for lang, translated_text in translations.items():
                            translated_subtitle = SubtitleEntry(
                                text=translated_text,
                                start_time=subtitle.start_time,
                                end_time=subtitle.end_time,
                                language=lang
                            )
                            self._add_to_subtitle_history(lang, translated_subtitle)

                # Create frames with appropriate subtitles for each language
                frames = {}
                for lang in self.languages:
                    # Get current subtitles to display
                    current_subtitles = self._get_current_subtitles(lang, current_time)

                    # Apply subtitles to frame
                    if current_subtitles:
                        subtitle_text = "\n".join([s.text for s in current_subtitles])
                        frames[lang] = self.subtitle_overlays[lang].apply(frame, subtitle_text)
                    else:
                        frames[lang] = frame.copy()

                # Write frames to RTMP outputs
                for lang, output_frame in frames.items():
                    await self.outputs[lang].write_frame_async(output_frame)

                # Update metrics
                frame_count += 1
                elapsed = time.time() - start_time
                self.metrics["fps"] = frame_count / elapsed
                self.metrics["end_to_end_latency"] = time.time() - frame_start_time

                # Adaptive quality based on performance
                self._adapt_to_performance()

                # Periodically log metrics
                if frame_count % 100 == 0:
                    self._log_metrics(frame_count, elapsed)

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

        # Clean up
        await self._cleanup()
        logger.info(f"Processed {frame_count} frames in {time.time() - start_time:.2f}s")

    def _add_to_subtitle_history(self, language: str, subtitle: SubtitleEntry):
        """Add a subtitle to the history for a specific language."""
        self.subtitle_history[language].append(subtitle)

        # Trim history if needed
        if len(self.subtitle_history[language]) > self.subtitle_history_length:
            self.subtitle_history[language].pop(0)

    def _get_current_subtitles(self, language: str, current_time: float) -> List[SubtitleEntry]:
        """Get subtitles that should be displayed at the current time."""
        return [
            s for s in self.subtitle_history[language]
            if s.start_time <= current_time <= s.end_time
        ]

    def _adapt_to_performance(self):
        """Adapt processing based on performance metrics."""
        # Example adaptive behavior
        if self.metrics["end_to_end_latency"] > 0.1:  # 100ms threshold
            # In a real implementation, this could reduce quality or processing
            pass

    def _log_metrics(self, frame_count: int, elapsed: float):
        """Log current performance metrics."""
        logger.info(
            f"Processed {frame_count} frames in {elapsed:.2f}s | "
            f"FPS: {self.metrics['fps']:.2f} | "
            f"Latency: {self.metrics['end_to_end_latency']*1000:.1f}ms | "
            f"Transcription: {self.metrics['transcription_latency']*1000:.1f}ms | "
            f"Translation: {self.metrics['translation_latency']*1000:.1f}ms"
        )

    async def _cleanup(self):
        """Clean up resources."""
        for output in self.outputs.values():
            await output.close()
        await self.stream.close()

    async def _mock_rtmp_stream(self) -> AsyncGenerator[Tuple[np.ndarray, np.ndarray], None]:
        """Mock RTMP stream for example purposes.
        In a real implementation, this would be replaced by RTMPCapture.aiter_frames_with_audio.
        """
        # Create a simple test video
        width, height = 1280, 720
        fps = 30
        duration = 60  # seconds

        # Create a mock audio waveform (silence with occasional "speech")
        sample_rate = 16000
        audio_frame_length = int(sample_rate / fps)  # audio samples per video frame

        # Generate random "speech segments"
        speech_segments = []
        t = 0
        while t < duration:
            segment_length = np.random.uniform(1.0, 4.0)  # 1-4 seconds of speech
            silence_length = np.random.uniform(2.0, 6.0)  # 2-6 seconds of silence
            speech_segments.append((t, t + segment_length))
            t += segment_length + silence_length

        # Simulate stream
        for i in range(int(fps * duration)):
            # Create frame with timestamp
            frame = np.zeros((height, width, 3), dtype=np.uint8)

            # Add a moving object to simulate video content
            x_pos = int((i % width) * (width - 100) / width)
            cv2.rectangle(frame, (x_pos, height // 2 - 50), (x_pos + 100, height // 2 + 50), (0, 0, 255), -1)

            # Add timestamp
            time_text = f"Time: {i/fps:.2f}s"
            cv2.putText(frame, time_text, (50, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2)

            # Create audio chunk (mostly silence)
            audio = np.zeros(audio_frame_length, dtype=np.float32)

            # Check if current time is in a speech segment
            current_time = i / fps
            in_speech = any(start <= current_time <= end for start, end in speech_segments)

            # If in speech segment, yield a random "transcription" with next frame
            if in_speech:
                # Add some "speech-like" noise to the audio
                audio = np.random.normal(0, 0.1, audio_frame_length)

                # For example purposes, we'll mock the transcription/translation
                # In a real implementation, this would be handled by the WhisperTranscriber
                if i % 30 == 0:  # New "utterance" every second
                    # Mock a transcription by adding text to the frame
                    phrases = [
                        "Hello, welcome to the livestream",
                        "Today we are discussing video processing",
                        "FMUS-VID makes it easy to work with videos",
                        "Real-time transcription is powerful",
                        "Let me show you this example"
                    ]
                    phrase = phrases[i // 30 % len(phrases)]

                    # In a real implementation, the transcriber would produce this text
                    # from the audio, and we'd inject it into the subtitle history
                    mock_subtitle = SubtitleEntry(
                        text=phrase,
                        start_time=current_time,
                        end_time=current_time + 3.0,
                        language="en"
                    )
                    self._add_to_subtitle_history("en", mock_subtitle)

                    # Mock translations
                    translations = {
                        "es": f"ES: {phrase}",
                        "fr": f"FR: {phrase}",
                        "de": f"DE: {phrase}"
                    }

                    for lang, text in translations.items():
                        mock_translation = SubtitleEntry(
                            text=text,
                            start_time=current_time,
                            end_time=current_time + 3.0,
                            language=lang
                        )
                        self._add_to_subtitle_history(lang, mock_translation)

            yield frame, audio

            # Simulate frame rate
            await asyncio.sleep(1/fps)

async def main_async():
    """Main async function to run the example."""
    # Configuration
    input_rtmp = "rtmp://localhost/live/stream"
    output_base_rtmp = "rtmp://localhost/live/output"
    languages = ["en", "es", "fr", "de"]

    # Create and run processor
    processor = RTMPSubtitleProcessor(
        input_rtmp=input_rtmp,
        output_base_rtmp=output_base_rtmp,
        languages=languages,
        whisper_model="medium",
        llm_model="ollama:mistral",
        device="cuda"
    )

    try:
        await processor.process()
    except KeyboardInterrupt:
        logger.info("Processing stopped by user")
    except Exception as e:
        logger.error(f"Error processing stream: {e}")

def main():
    """Run the RTMP subtitle example."""
    logger.info("Starting RTMP multilingual subtitle example")

    # In a real implementation, we would check for required components
    logger.info("Note: This is a conceptual example. The required modules must be implemented.")
    logger.info("Required modules: RTMPCapture, RTMPOutput, WhisperTranscriber, LLMTranslator")

    # Show what this would do in a real implementation
    logger.info("\nThis example would:")
    logger.info("1. Capture an RTMP stream from: rtmp://localhost/live/stream")
    logger.info("2. Transcribe speech using FasterWhisper")
    logger.info("3. Translate subtitles to Spanish, French, and German")
    logger.info("4. Output streams with different subtitles to:")
    logger.info("   - rtmp://localhost/live/output_en (English)")
    logger.info("   - rtmp://localhost/live/output_es (Spanish)")
    logger.info("   - rtmp://localhost/live/output_fr (French)")
    logger.info("   - rtmp://localhost/live/output_de (German)")

    # Run the async pipeline
    asyncio.run(main_async())

if __name__ == "__main__":
    main()
