stream_state.py

from __future__ import annotations

import threading
from dataclasses import dataclass
from app_types import SegmentPlayMode
from segments import Segment, Segments
from util import *


@dataclass
class StreamStateSnapshot:
    """
    Immutable snapshot of stream state for use in audio callback.
    Created by StreamState.begin_playback().
    """
    index: int
    paused: bool
    min_index: int
    max_index: int
    segments: list[Segment]
    segment_play_mode: SegmentPlayMode | None
    compressor_strength: float
    compressor_enabled: bool
    sample_rate: int


class StreamState:
    """
    Thread-safe container for various `Stream` state values 
    (but not including the actual audio data).
    
    Any logic that relies on multiple reads/writes of these values
    are/should be encapsulated in functions that uses a single lock operation.
    """

    def __init__(self):
        self._lock = threading.Lock()

        self._index: int = 0
        self._paused: bool = True
        self._min_index: int = 0
        self._max_index: int = 0
        self._segments: list[Segment] = []
        self._segment_play_mode: SegmentPlayMode | None = None
        self._compressor_strength: float = DEFAULT_COMPRESSOR_STRENGTH
        self._compressor_enabled: bool = False
        self._sample_rate: int = 0

    # --- Snapshot API for audio callback ---

    def begin_playback(self) -> StreamStateSnapshot:
        """
        Called by audio callback at start of each frame.
        Returns an immutable snapshot of all state needed for playback.
        """
        with self._lock:
            return StreamStateSnapshot(
                index=self._index,
                paused=self._paused,
                min_index=self._min_index,
                max_index=self._max_index,
                segments=self._segments.copy() if self._segments else [],
                segment_play_mode=self._segment_play_mode,
                compressor_strength=self._compressor_strength,
                compressor_enabled=self._compressor_enabled,
                sample_rate=self._sample_rate
            )

    def end_playback(self, new_index: int) -> None:
        """
        Called by audio callback at end of each frame.
        Writes back the updated index.
        """
        with self._lock:
            self._index = new_index

    # --- Individual property access for UI thread ---

    @property
    def paused(self) -> bool:
        with self._lock:
            return self._paused

    @paused.setter
    def paused(self, value: bool) -> None:
        with self._lock:
            self._paused = value

    @property
    def index(self) -> int:
        with self._lock:
            return self._index

    @index.setter
    def index(self, value: int) -> None:
        with self._lock:
            self._index = value

    @property
    def min_index(self) -> int:
        with self._lock:
            return self._min_index

    @min_index.setter
    def min_index(self, value: int) -> None:
        with self._lock:
            self._min_index = value

    @property
    def max_index(self) -> int:
        with self._lock:
            return self._max_index

    @max_index.setter
    def max_index(self, value: int) -> None:
        with self._lock:
            self._max_index = value

    @property
    def segments(self) -> list[Segment]:
        with self._lock:
            return self._segments

    @segments.setter
    def segments(self, value: list[Segment]) -> None:
        with self._lock:
            self._segments = value
            # Reset segment play mode if segments cleared
            if not self._segments and self._segment_play_mode is not None:
                self._segment_play_mode = None

    @property
    def segment_play_mode(self) -> SegmentPlayMode | None:
        with self._lock:
            return self._segment_play_mode

    @segment_play_mode.setter
    def segment_play_mode(self, value: SegmentPlayMode | None) -> None:
        with self._lock:
            self._segment_play_mode = value

    @property
    def compressor_strength(self) -> float:
        with self._lock:
            return self._compressor_strength

    @compressor_strength.setter
    def compressor_strength(self, value: float) -> None:
        with self._lock:
            self._compressor_strength = value

    @property
    def compressor_enabled(self) -> bool:
        with self._lock:
            return self._compressor_enabled

    @compressor_enabled.setter
    def compressor_enabled(self, value: bool) -> None:
        with self._lock:
            self._compressor_enabled = value

    # --- Batch operations ---

    def reset(self, max_index: int) -> None:
        """
        Reset state for a new sound.
        Sets min_index to 0, max_index to given value, index to 0, clears segments.
        """
        with self._lock:
            self._min_index = 0
            self._max_index = max_index
            self._index = 0
            self._segments = []

    def clear(self) -> None:
        """
        Clear all state.
        """
        with self._lock:
            self._index = 0
            self._min_index = 0
            self._max_index = 0
            self._segments = []
            self._segment_play_mode = None
            self._sample_rate = 0

    def set_index_clamped(self, value: int) -> None:
        """
        Sets index, clamping to min/max range.
        """
        with self._lock:
            self._index = max(self._min_index, min(value, self._max_index))

    def set_min_max_index(self, min_index: int, max_index: int) -> None:
        """
        Sets min/max indices
        Clamp current index to the new range.
        """
        if max_index <= min_index:
            return
        with self._lock:
            self._min_index = min_index
            self._max_index = max_index
            self._index = clamp(self._index, self._min_index, self._max_index)

    def set_min_max_index_using_positions(
        self, min_max_position: tuple[float, float], data_len: int
    ) -> None:
        with self._lock:
            min_position, max_position = min_max_position
            min_index = int(min_position * self._sample_rate)
            min_index = clamp(min_index, 0.0, data_len - 1)
            max_index = int(max_position * self._sample_rate)
            max_index = clamp(max_index, 0.0, data_len - 1)
            if max_index <= min_index:
                return
        self.set_min_max_index(min_index, max_index)

    @property
    def sample_rate(self) -> int:
        with self._lock:
            return self._sample_rate

    @sample_rate.setter
    def sample_rate(self, value: int) -> None:
        with self._lock:
            self._sample_rate = value

    # --- Computed properties (derived from state) ---

    @property
    def position(self) -> float:
        """Gets the position in seconds (derived from index)."""
        with self._lock:
            if self._sample_rate == 0:
                return 0
            return self._index / self._sample_rate

    @position.setter
    def position(self, seconds: float) -> None:
        """Sets position in seconds (converted to index)."""
        with self._lock:
            if self._sample_rate == 0:
                return
            index = int(seconds * self._sample_rate)
            self._index = max(self._min_index, min(index, self._max_index))

    @property
    def min_position(self) -> float:
        """Gets the start position in seconds."""
        with self._lock:
            if self._sample_rate == 0:
                return 0
            return self._min_index / self._sample_rate

    @property
    def max_position(self) -> float:
        """Gets the end position in seconds."""
        with self._lock:
            if self._sample_rate == 0:
                return 0
            return self._max_index / self._sample_rate

    @property
    def is_at_end(self) -> bool:
        """Checks if playback has reached the end."""
        with self._lock:
            if self._segment_play_mode == SegmentPlayMode.TO_END:
                if not self._segments:
                    return False
                max_position = self._segments[-1].end
                max_index = int(max_position * self._sample_rate) if self._sample_rate else 0
                return self._index >= max_index
            else:
                return self._index >= self._max_index

    def set_position_by_segment_index(self, index: int) -> None:
        """Sets position to the start of a specific segment."""
        with self._lock:
            if not self._segments or index > len(self._segments) - 1:
                return
            if self._sample_rate == 0:
                return
            self._index = int(self._segments[index].start * self._sample_rate) + 1

    def get_segment_index(self) -> int:
        """Gets the current segment index based on position."""
        with self._lock:
            if self._sample_rate == 0:
                return -1
            position = self._index / self._sample_rate
            return Segments.segment_index_of(position, self._segments)


DEFAULT_COMPRESSOR_STRENGTH = 2.5 # TODO: ui, parameterize, etc