stream.py

from __future__ import annotations

import numpy as np
import queue
import sounddevice as sd
import threading
from numpy import ndarray
from typing import Optional
from app_types import SegmentPlayMode
from segments import Segment, Segments
from simple_compressor import SimpleCompressor
from sound import Sound
from sound_util import SoundUtil
from stream_state import StreamState
from sys_sound_util import SysSoundUtil
from util import *


class Stream:
    """
    Streams audio to the default audio output device.
    
    State management is delegated to StreamState, which handles thread-safety
    for attributes accessed from both UI and audio callback threads.
    """

    def __init__(self):
        """
        Initializes the AudioStreamer.
        """
        
        # Thread-safe queue for signaling stream completion from audio callback
        # Applies to non-segment-mode only
        self.completion_queue = queue.Queue()

        # The data buffer from which the audio stream callback function draws from to
        # stream the audio in realtime
        self._data = np.array([], dtype=np.float32).reshape(0, 2)
       
        # Playback state (thread-safe via StreamState's internal lock)
        self._state = StreamState()
        
        self._stop_requested = threading.Event()
        
        # Tracks whether completion has been signaled for current playback
        # Prevents multiple completion events from being queued
        self._completion_signaled = threading.Event()

        # The sounddevice stream object (is None until start() is called)
        self._stream: Optional[sd.OutputStream] = None

        # Audio compressor for dynamic range compression
        self._compressor = SimpleCompressor()

    # --- Public state access (delegates to StreamState) ---

    @property
    def paused(self) -> bool:
        return self._state.paused

    def pause(self) -> None:
        """
        Pauses the audio stream. While paused, the stream will output silence.
        """
        self._state.paused = True

    def unpause(self) -> None:
        """
        Resumes the audio stream after it has been paused.
        """
        self._state.paused = False

    @property
    def position(self) -> float:
        """Gets the position in seconds."""
        return self._state.position
    
    @position.setter
    def position(self, seconds: float) -> None:
        self._state.position = seconds

    def set_position_by_segment_index(self, index: int) -> None:
        """Sets position to the start of a specific segment."""
        self._state.set_position_by_segment_index(index)

    @property
    def min_position(self) -> float:
        """Gets the start position in seconds."""
        return self._state.min_position

    @property
    def max_position(self) -> float:
        return self._state.max_position
        
    @property
    def is_at_end(self) -> bool:
        return self._state.is_at_end
    
    def set_min_max_position(self, min_max_position: tuple[float, float]) -> None:
        """ Sets min_index and max_index using time values """
        self._state.set_min_max_index_using_positions(min_max_position, len(self._data))

    @property
    def segments(self) -> list[Segment]:
        return self._state.segments
    
    @segments.setter
    def segments(self, value: list[Segment]) -> None:
        self._state.segments = value

    def get_segment_index(self) -> int:
        """Gets the current segment index based on position."""
        return self._state.get_segment_index()

    def get_completion_event(self) -> bool:
        """
        Check if a stream completion event has occurred.
        Non-blocking check of the completion queue.
        Returns True if stream completed, False otherwise.
        """
        try:
            return self.completion_queue.get_nowait()
        except queue.Empty:
            return False

    def clear_completion_event(self) -> None:
        """Clear the completion signal flag to allow future completion events."""
        self._completion_signaled.clear()

    @property
    def segment_play_mode(self) -> SegmentPlayMode | None:
        return self._state.segment_play_mode
    
    @segment_play_mode.setter
    def segment_play_mode(self, value: SegmentPlayMode | None) -> None:
        self._state.segment_play_mode = value

    @property
    def compressor_enabled(self) -> bool:
        """Check if audio compression is enabled."""
        return self._state.compressor_enabled

    @compressor_enabled.setter
    def compressor_enabled(self, value: bool) -> None:
        """Enable or disable audio compression."""
        self._state.compressor_enabled = value
        if value:
            self._compressor.reset()  # Clear envelope when enabling

    @property
    def compressor_strength(self) -> float:
        """Get the compressor strength."""
        return self._state.compressor_strength

    @compressor_strength.setter
    def compressor_strength(self, value: float) -> None:
        """Set the compressor strength."""
        self._state.compressor_strength = value

    def _init_and_start_stream(self, sr: int) -> None:
        """
        Creates and starts the output stream.
        Called internally by start() and when sample rate changes.
        """
        self._state.sample_rate = sr
        self._compressor.set_sample_rate(sr)

        BASE_BLOCKSIZE = 128
        if sr < 40000: # eg, 22050, 24000
            blocksize = BASE_BLOCKSIZE * 1
        elif sr < 80000: # eg, 44100, 48000
            blocksize = BASE_BLOCKSIZE * 2
        elif sr < 160000: # eg, 88200, 96000
            blocksize = BASE_BLOCKSIZE * 4
        else: # eg, 192k, ...
            blocksize = BASE_BLOCKSIZE * 8

        self._stream = sd.OutputStream(
            samplerate=sr,
            channels=2,
            callback=self._callback,
            dtype=np.float32,  # We work with float32 internally
            blocksize=blocksize,
            latency="low"
        )
        try:
            self._stream.start()
        except Exception as e:
            print("*", e)

    def _kill_stream(self) -> None:
        """
        Stops and closes the audio stream, releasing all resources.
        Also clears data buffer.
        """
        if self._stream:
            self._stream.stop()
            self._stream.close()
            self._stream = None
        self._state.sample_rate = 0

    def init_sound(self, sound: Sound) -> None:
        """
        Inits stream (using sound's samplerate if possible),
        and sets sound buffer using sound's data.
        
        Kills and restarts stream if if samplerate has changed.

        ---
        Rem, system's audio layer may resample audio regardless:

          If using pipewire (linux), configure it to accept the gamut of typical sample rates,
          and then run "pw-top" to verify. If pipewire and sound device both support the source sample rate,
          pretty sure unnecessary resampling is avoided.

          If using Windows, cannot be defeated without using exclusive mode (WASAPI/ASIO).
        """

        supported_sample_rates = SysSoundUtil.get_supported_sample_rates_of_device()

        if sound.sr in supported_sample_rates:
            target_sr = sound.sr
            data = sound.data
        else:
            target_sr = SysSoundUtil.get_current_device_default_sample_rate()
            data = SoundUtil.resample(sound, target_sr).data
            print(f"Stream will use device default sample rate {target_sr} because source sound's sample rate {sound.sr} not supported by output device {supported_sample_rates}")

        if self._stream and self._state.sample_rate != target_sr:
            self._kill_stream()

        if not self._stream:
            print("Starting stream, sr", target_sr)
            self._init_and_start_stream(target_sr)

        self._data = self._convert_data_for_stream(data)

        # Reset state for new sound
        self._state.reset(len(self._data) - 1)
        
        # Reset completion signaling for new playback
        self._completion_signaled.clear()

    def clear(self) -> None:
        self._data = np.array([], dtype=np.float32).reshape(0, 2)
        self._state.clear()

    def _convert_data_for_stream(self, data: ndarray) -> ndarray:
        """
        Converts incoming data to stereo samples-last format (n_samples, 2) for the output stream.
        Does not resample.
        
        Input formats accepted:
        - Mono 1D: (n_samples,) -> stereo by duplicating to both channels
        - Stereo 2D channels-first: (2, n_samples) -> transpose to (n_samples, 2)
        - Stereo 2D samples-last: (n_samples, 2) -> pass through
        """
        # Handle integer to float conversion first
        if np.issubdtype(data.dtype, np.integer):
            max_val = np.iinfo(data.dtype).max
            data = data.astype(np.float32) / max_val
        elif data.dtype != np.float32:
            data = data.astype(np.float32)
        
        # Convert to stereo samples-last format (n_samples, 2)
        if data.ndim == 1:
            # Mono: duplicate to both channels -> (n_samples, 2)
            data = np.column_stack((data, data))
        elif data.ndim == 2:
            if data.shape[0] == 2 and data.shape[1] != 2:
                # Channels-first stereo (2, n_samples) -> transpose to (n_samples, 2)
                data = data.T
            elif data.shape[1] == 2:
                # Already samples-last stereo (n_samples, 2) -> pass through
                pass
            else:
                # Unexpected shape, try to make sense of it
                # If first dim is small, assume channels-first
                if data.shape[0] < data.shape[1]:
                    data = data.T
                else:
                    # Fallback: take first two columns or duplicate
                    if data.shape[1] >= 2:
                        data = data[:, :2]
                    else:
                        data = np.column_stack((data[:, 0], data[:, 0]))
        else:
            raise ValueError(f"Unsupported audio data dimensions: {data.ndim}")
        
        return data

    def _callback(self, outdata: ndarray, frames: int, time, status: sd.CallbackFlags) -> None:
        """
        The heart of the audio streamer, called by sounddevice in a separate thread.
        It pulls data from the buffer and sends it to the audio output.
        
        outdata shape is (frames, 2) for stereo output.
        """

        if self._stop_requested.is_set():
            outdata.fill(0)
            raise sd.CallbackStop

        if status.output_underflow:
            print("warning: Output underflow")

        # Get snapshot of state at start of callback
        state = self._state.begin_playback()
        
        # Work with local variables from snapshot
        if state.segment_play_mode is not None and state.segments:
            new_index = self._callback_segments_mode(outdata, frames, state)
        else:
            new_index = self._callback_normal(outdata, frames, state)
        
        # Write back updated index
        self._state.end_playback(new_index)
            
    def _callback_normal(self, outdata: ndarray, frames: int, state) -> int:
        """
        Normal (non-segment-mode) playback callback.
        Returns the new index value.
        """

        if state.paused:
            # If paused, stream silence
            outdata.fill(0)
            return state.index
        
        index = state.index
        output_pos = 0
        remaining = frames

        while remaining > 0:
            # Calculate available samples from current cursor position
            available = state.max_index - index

            if available <= 0:
                # Reached end of playback range; fill rest with silence and signal complete
                if not self._completion_signaled.is_set():
                    self.completion_queue.put(True)
                    self._completion_signaled.set()
                outdata[output_pos:] = 0
                return index

            # Copy as much as we can from current position
            samples_to_copy = min(available, remaining)
            chunk = self._data[index:index + samples_to_copy].copy()
            if state.compressor_enabled:
                chunk = self._compressor.process(chunk, strength=state.compressor_strength)
            outdata[output_pos:output_pos + samples_to_copy] = chunk

            index += samples_to_copy
            output_pos += samples_to_copy
            remaining -= samples_to_copy
        
        return index

    def _clamp_segment_index(self, index: int, state) -> int:
        """
        Clamp index to a valid segment position using the same rules
        as _callback_segments_mode. Used when paused to ensure the
        returned index is always at a valid segment position.
        """
        sr = state.sample_rate
        current_time = index / sr if sr else 0

        # TO_END: clamp to last segment end if past it
        if state.segment_play_mode == SegmentPlayMode.TO_END:
            last_segment_end = state.segments[-1].end
            if current_time >= last_segment_end:
                return int(last_segment_end * sr)

        where_desc, i = Segments.where(current_time, state.segments)

        if where_desc == "inside":
            return index  # Already at a valid position

        # In a gap — apply same jump logic as the main loop
        dest_segment_index = None

        match where_desc:
            case "before_first":
                dest_segment_index = 0
            case "before":
                assert isinstance(i, int)
                if state.segment_play_mode == SegmentPlayMode.REPEAT_SINGLE:
                    dest_segment_index = i - 1
                else:
                    dest_segment_index = i
            case "after_last":
                if state.segment_play_mode == SegmentPlayMode.REPEAT_ALL:
                    dest_segment_index = 0
                elif state.segment_play_mode == SegmentPlayMode.REPEAT_SINGLE:
                    dest_segment_index = len(state.segments) - 1

        if dest_segment_index is not None:
            position = state.segments[dest_segment_index].start
            return int(position * sr) + 1

        return index

    def _callback_segments_mode(self, outdata: ndarray, frames: int, state) -> int:
        """
        Segment-aware playback:
        - Plays only from segment time ranges
        - Skips gaps by jumping index to next segment start
        - Outputs silence when past the last segment
        
        Returns the new index value.
        """

        if state.paused:
            # If paused, stream silence but still clamp the index
            outdata.fill(0)
            return self._clamp_segment_index(state.index, state)

        index = state.index
        output_pos = 0
        remaining = frames
        sr = state.sample_rate

        while remaining > 0:

            current_time = index / sr if sr else 0

            # Check if we've reached/passed end of last segment in SEGMENT_TO_END mode
            if state.segment_play_mode == SegmentPlayMode.TO_END:
                last_segment_end = state.segments[-1].end
                if current_time >= last_segment_end:
                    outdata[output_pos:] = 0
                    return index
            else:
                ...

            # Find which segment we're in (or next segment if in gap)
            where_desc, i = Segments.where(current_time, state.segments)
            segment_index = i if where_desc == "inside" else None
            next_segment_index = i if where_desc == "before" else None
           
            if segment_index is not None: # We're inside a segment

                segment_end_time = state.segments[segment_index].end
                seg_end_index = int(segment_end_time * sr)
                
                # Calculate how many samples we can play before segment ends
                samples_until_segment_end = max(0, seg_end_index - index)
                samples_to_play = min(remaining, samples_until_segment_end)
                
                if samples_to_play == 0:
                    # Reached segment end, increment index to move past segment
                    # so next iteration will trigger segment transition logic
                    index += 1
                    continue
                
                # Copy audio data to output
                chunk = self._data[index:index + samples_to_play].copy()
                if state.compressor_enabled:
                    chunk = self._compressor.process(chunk, strength=state.compressor_strength)
                outdata[output_pos:output_pos + samples_to_play] = chunk
                
                # Update position and remaining frames
                index += samples_to_play
                output_pos += samples_to_play
                remaining -= samples_to_play
            
            else: # We're outside any segment
                
                if state.segment_play_mode == SegmentPlayMode.TO_END and where_desc == "after_last":
                    # Clamp index and play silence
                    index = int(state.segments[-1][1] * sr)
                    outdata[output_pos:] = 0
                    return index

                dest_segment_index = None

                match where_desc:
                    case "before_first":
                        dest_segment_index = 0
                    case "before":
                        assert isinstance(next_segment_index, int)
                        if state.segment_play_mode == SegmentPlayMode.REPEAT_SINGLE:
                            dest_segment_index = next_segment_index - 1
                        else:
                            dest_segment_index = next_segment_index
                    case "after_last":
                        if state.segment_play_mode == SegmentPlayMode.REPEAT_ALL:
                            dest_segment_index = 0
                        elif state.segment_play_mode == SegmentPlayMode.REPEAT_SINGLE:
                            dest_segment_index = len(state.segments) - 1

                if dest_segment_index is None:
                    raise Exception("Logic error, dest segment index unassigned")

                # print("dest segment index", dest_segment_index)
                position = state.segments[dest_segment_index].start
                index = int(position * sr) + 1 # "+1" is important here
        
        return index