"""Streaming inference for SalTTS."""

import torch
import numpy as np
from typing import Generator, Optional
import time

from saltts.inference.engine import SalTTSEngine
from saltts.core.data_structures import AudioOutput


class StreamingInference:
    """Real-time streaming inference for SalTTS.

    Processes input in chunks for low-latency streaming.
    """

    def __init__(
        self,
        engine: SalTTSEngine,
        chunk_size: int = 3,
        max_latency_ms: float = 300.0,
    ):
        """Initialize streaming inference.

        Args:
            engine: SalTTS engine
            chunk_size: Chunk size (tokens)
            max_latency_ms: Maximum latency (ms)
        """
        self.engine = engine
        self.chunk_size = chunk_size
        self.max_latency_ms = max_latency_ms

    def stream_synthesize(
        self,
        text: str,
        emotion: str = "neutral",
        speaker_id: int = 0,
    ) -> Generator[np.ndarray, None, None]:
        """Stream synthesize speech from text.

        Args:
            text: Input text
            emotion: Emotion label
            speaker_id: Speaker ID

        Yields:
            Audio chunks
        """
        # Split text into chunks
        words = text.split()
        for i in range(0, len(words), self.chunk_size):
            chunk_text = " ".join(words[i : i + self.chunk_size])

            # Synthesize chunk
            audio_output = self.engine.synthesize(
                chunk_text,
                emotion=emotion,
                speaker_id=speaker_id,
            )

            yield audio_output.waveform
