#!/usr/bin/env python3
"""
Smart Progressive Streaming Handler

Provides frequent partial transcriptions (every 500ms) with:
- Growing window up to 15s for accuracy
- Sentence-boundary-aware window sliding for audio > 15s
- Fixed sentences + active transcription
"""

from dataclasses import dataclass
from types import SimpleNamespace
from typing import Any, Generator

import numpy as np


@dataclass
class PartialTranscription:
    """Result from progressive streaming."""

    fixed_text: str  # Sentences that won't change
    active_text: str  # Current partial transcription
    timestamp: float  # Current position in audio
    is_final: bool  # True if this is the last update


class SmartProgressiveStreamingHandler:
    """
    Smart progressive streaming with sentence-aware window management.

    Strategy:
    1. Emit partial transcriptions every 500ms
    2. Use growing window (up to 15s) for better accuracy
    3. When audio > 15s, slide window using sentence boundaries:
       - Keep completed sentences as "fixed"
       - Only re-transcribe the "active" portion
    """

    def __init__(
        self,
        model: Any,
        emission_interval: float = 0.5,
        max_window_size: float = 15.0,
        sentence_buffer: float = 2.0,
    ) -> None:
        """
        Args:
            model: Parakeet model with sentence alignment
            emission_interval: Emit partial transcription every N seconds (default 500ms)
            max_window_size: Maximum window size before sliding (default 15s)
            sentence_buffer: Keep last N seconds of sentences in active window (default 2s)
        """
        self.model = model
        self.emission_interval = emission_interval
        self.max_window_size = max_window_size
        self.sentence_buffer = sentence_buffer
        self.sample_rate = 16000

        # State for incremental streaming
        self.reset()

    def reset(self) -> None:
        """Reset state for new streaming session."""
        self.fixed_sentences: list[str] = []
        self.fixed_end_time: float = 0.0
        self.last_transcribed_length: int = 0

    def _decode_window(self, audio_window: np.ndarray) -> Any:
        """
        Decode an audio window and return an object with:
          - text: full transcript for the window
          - sentences: list of objects with .text and .end (seconds)
        """
        if hasattr(self.model, "decode_chunk"):
            import mlx.core as mx

            audio_mx = mx.array(audio_window, dtype=mx.float32)
            return self.model.decode_chunk(audio_mx, verbose=False)

        # nano-parakeet: use timestamps from transcribe
        result = self.model.transcribe(audio_window, timestamps=True)
        segments = result.timestamp.get("segment", []) if hasattr(result, "timestamp") else []
        sentences = [SimpleNamespace(text=seg.get("segment", "").strip(), end=seg.get("end", 0.0)) for seg in segments]
        return SimpleNamespace(text=getattr(result, "text", ""), sentences=sentences)

    def transcribe_incremental(self, audio: np.ndarray) -> PartialTranscription:
        """
        Transcribe audio incrementally (for live streaming).

        Call this repeatedly with growing audio buffer.
        Returns a single PartialTranscription for current state.
        """
        # Skip if not enough new audio
        current_length = len(audio)
        if current_length < self.sample_rate * 0.5:  # Need at least 500ms
            return PartialTranscription(
                fixed_text=" ".join(self.fixed_sentences),
                active_text="",
                timestamp=current_length / self.sample_rate,
                is_final=False,
            )

        # Skip if no new audio since last transcription
        if current_length == self.last_transcribed_length:
            return PartialTranscription(
                fixed_text=" ".join(self.fixed_sentences),
                active_text="",
                timestamp=current_length / self.sample_rate,
                is_final=False,
            )

        self.last_transcribed_length = current_length

        # Extract window for transcription (from last fixed sentence to end)
        window_start_samples = int(self.fixed_end_time * self.sample_rate)
        audio_window = audio[window_start_samples:]

        # Transcribe current window
        result = self._decode_window(audio_window)

        # Check if window exceeds max_window_size
        window_duration = len(audio_window) / self.sample_rate

        if window_duration >= self.max_window_size and len(result.sentences) > 1:
            # Window is too large - fix some sentences
            cutoff_time = window_duration - self.sentence_buffer

            # Find sentences to fix
            new_fixed_sentences = []
            new_fixed_end_time = self.fixed_end_time

            for sentence in result.sentences:
                sentence_abs_time = self.fixed_end_time + sentence.end

                if sentence.end < cutoff_time:
                    # Fix this sentence
                    new_fixed_sentences.append(sentence.text.strip())
                    new_fixed_end_time = sentence_abs_time
                else:
                    break

            if new_fixed_sentences:
                self.fixed_sentences.extend(new_fixed_sentences)
                self.fixed_end_time = new_fixed_end_time

                # Re-transcribe from new fixed point
                window_start_samples = int(self.fixed_end_time * self.sample_rate)
                audio_window = audio[window_start_samples:]
                result = self._decode_window(audio_window)

        # Build output
        fixed_text = " ".join(self.fixed_sentences)
        active_text = result.text.strip()
        timestamp = len(audio) / self.sample_rate

        return PartialTranscription(fixed_text=fixed_text, active_text=active_text, timestamp=timestamp, is_final=False)

    def transcribe_progressive(self, audio: np.ndarray) -> Generator[PartialTranscription, None, None]:
        """
        Transcribe audio with smart progressive emissions.

        Yields PartialTranscription with:
        - fixed_text: Completed sentences (won't change)
        - active_text: Current partial transcription
        - timestamp: Current position
        - is_final: True for last update
        """
        position = 0  # Start of current window (in samples)
        fixed_sentences = []  # List of completed sentence texts
        fixed_end_time = 0.0  # End time of last fixed sentence

        while position < len(audio):
            # Determine current window end
            remaining = (len(audio) - position) / self.sample_rate

            if remaining <= self.emission_interval:
                # Last chunk - process everything remaining
                window_end = len(audio)
                is_final = True
            else:
                # Regular chunk
                window_end = min(len(audio), position + int(self.emission_interval * self.sample_rate))
                is_final = False

            # Extract window for transcription
            # Window includes: fixed_end_time to current position
            # (we re-transcribe from last fixed sentence to get better context)
            window_start_samples = int(fixed_end_time * self.sample_rate)
            audio_window = audio[window_start_samples:window_end]

            # Transcribe current window
            result = self._decode_window(audio_window)

            # Check if window exceeds max_window_size
            window_duration = (window_end - window_start_samples) / self.sample_rate

            if window_duration >= self.max_window_size and len(result.sentences) > 1:
                # Window is too large - need to fix some sentences
                # Strategy: Keep sentences up to (max_window - sentence_buffer) as fixed
                cutoff_time = window_duration - self.sentence_buffer

                # Find sentences that are completed and before cutoff
                new_fixed_sentences = []
                new_fixed_end_time = fixed_end_time

                for sentence in result.sentences:
                    sentence_abs_time = fixed_end_time + sentence.end

                    if sentence.end < cutoff_time:
                        # This sentence is complete and old enough to fix
                        new_fixed_sentences.append(sentence.text.strip())
                        new_fixed_end_time = sentence_abs_time
                    else:
                        # Stop here - keep remaining sentences as active
                        break

                if new_fixed_sentences:
                    fixed_sentences.extend(new_fixed_sentences)
                    fixed_end_time = new_fixed_end_time

                    # Re-transcribe from new fixed_end_time
                    window_start_samples = int(fixed_end_time * self.sample_rate)
                    audio_window = audio[window_start_samples:window_end]
                    result = self._decode_window(audio_window)

            # Build output
            fixed_text = " ".join(fixed_sentences)
            active_text = result.text.strip()
            timestamp = window_end / self.sample_rate

            yield PartialTranscription(
                fixed_text=fixed_text, active_text=active_text, timestamp=timestamp, is_final=is_final
            )

            # Move to next position
            position = window_end

            if is_final:
                break


def demo_smart_progressive() -> None:
    """Demonstrate smart progressive streaming."""
    import time

    import soundfile as sf
    from mlx_audio.stt.generate import load_model

    print("=" * 80)
    print("SMART PROGRESSIVE STREAMING DEMO")
    print("=" * 80)
    print("\nFeatures:")
    print("  • Partial updates every 500ms")
    print("  • Growing window up to 15s")
    print("  • Sentence-aware window sliding for long audio")
    print()

    # Load model
    print("Loading model...")
    model = load_model("mlx-community/parakeet-tdt-0.6b-v3")

    # Load audio
    audio, sr = sf.read("reachy-voice-test.wav")
    if len(audio.shape) > 1:
        audio = audio.mean(axis=1)
    if sr != 16000:
        from scipy import signal

        audio = signal.resample(audio, int(len(audio) * 16000 / sr))
    audio = audio.astype(np.float32)

    # Test cases
    test_cases = [
        (audio, "Short (5.8s)"),
        (np.concatenate([audio] * 3), "Medium (17.5s) - exceeds 15s window"),
        (np.concatenate([audio] * 5), "Long (29.1s) - multiple sentence fixings"),
    ]

    for test_audio, label in test_cases:
        duration = len(test_audio) / 16000

        print(f"\n{'=' * 80}")
        print(f"TEST: {label} - Duration: {duration:.1f}s")
        print("=" * 80)

        handler = SmartProgressiveStreamingHandler(
            model,
            emission_interval=0.5,  # 500ms updates
            max_window_size=15.0,  # Max 15s window
            sentence_buffer=2.0,  # Keep 2s of sentences in active window
        )

        print("\nProgressive transcriptions (showing every 4th update for readability):")
        print()

        update_count = 0
        start_time = time.perf_counter()

        for result in handler.transcribe_progressive(test_audio):
            update_count += 1

            # Show every other update (every 1s) to keep output manageable
            if update_count % 2 == 0 or result.is_final:
                elapsed = time.perf_counter() - start_time
                marker = "FINAL" if result.is_final else f"Update {update_count}"

                print(f"[{result.timestamp:5.1f}s | {elapsed:.3f}s elapsed] {marker}:")

                if result.fixed_text:
                    print(f"  Fixed:  {result.fixed_text[:70]}...")
                print(f"  Active: {result.active_text[:70]}...")
                print()

        total_time = time.perf_counter() - start_time
        print(f"Total updates: {update_count}")
        print(f"Total time: {total_time:.3f}s ({total_time / duration:.3f}s per audio second)")

    print("\n" + "=" * 80)
    print("KEY INSIGHTS")
    print("=" * 80)
    print("""
1. Frequent updates (500ms):
   - User sees transcription building in real-time
   - Low latency feel

2. Growing window (up to 15s):
   - Better accuracy with more context
   - No arbitrary chunking for short audio

3. Sentence-aware sliding (for audio > 15s):
   - Completed sentences become "fixed"
   - Only active portion re-transcribed
   - Maintains accuracy while managing window size

4. Performance:
   - Processing happens DURING user speaking
   - Final result available immediately when user stops
    """)


if __name__ == "__main__":
    demo_smart_progressive()
