src/state.py

"""
User state management for the AIGC Telegram Bot.
Tracks per-user conversation state and processing tasks.
"""

import time
from dataclasses import dataclass, field
from enum import Enum
from typing import Dict, List, Optional, Any
from collections import defaultdict
import threading


class UserState(Enum):
    """States a user can be in during a conversation."""
    IDLE = "idle"
    AWAITING_SONG_INPUT = "awaiting_song_input"
    AWAITING_VOICE_MODEL = "awaiting_voice_model"
    AWAITING_YOUTUBE_URL = "awaiting_youtube_url"
    PROCESSING = "processing"
    AWAITING_SETTINGS = "awaiting_settings"
    AWAITING_PITCH_CHANGE = "awaiting_pitch_change"


class InputType(Enum):
    """Type of song input."""
    YOUTUBE = "yt"
    FILE = "local"


@dataclass
class UserSession:
    """Stores per-user session data."""
    state: UserState = UserState.IDLE
    input_type: Optional[InputType] = None

    # Song input
    song_input: Optional[str] = None  # URL or file path
    song_file_id: Optional[str] = None  # Telegram file_id for caching

    # Voice model selection
    voice_model: Optional[str] = None

    # RVC parameters
    pitch_change: int = 0
    index_rate: float = 0.5
    filter_radius: int = 3
    volume_envelope: float = 0.25
    f0_method: str = "rmvpe"
    hop_length: int = 128
    protect: float = 0.33

    # Audio effects
    reverb_size: float = 0.15
    reverb_wet: float = 0.2
    reverb_dry: float = 0.8
    reverb_damping: float = 0.7

    # Mix gains
    main_gain: int = 0
    backup_gain: int = 0
    inst_gain: int = 0

    # Output settings
    output_format: str = "mp3"
    inference_mode: str = "full"
    keep_files: bool = False
    extra_denoise: bool = False
    pitch_change_all: int = 0

    # Metadata
    created_at: float = field(default_factory=time.time)
    last_activity: float = field(default_factory=time.time)

    def to_dict(self) -> Dict[str, Any]:
        """Serialize session to dict."""
        return {
            "state": self.state.value,
            "input_type": self.input_type.value if self.input_type else None,
            "voice_model": self.voice_model,
            "pitch_change": self.pitch_change,
            "index_rate": self.index_rate,
            "filter_radius": self.filter_radius,
            "volume_envelope": self.volume_envelope,
            "f0_method": self.f0_method,
            "hop_length": self.hop_length,
            "protect": self.protect,
            "reverb_size": self.reverb_size,
            "reverb_wet": self.reverb_wet,
            "reverb_dry": self.reverb_dry,
            "reverb_damping": self.reverb_damping,
            "main_gain": self.main_gain,
            "backup_gain": self.backup_gain,
            "inst_gain": self.inst_gain,
            "output_format": self.output_format,
            "inference_mode": self.inference_mode,
            "keep_files": self.keep_files,
            "extra_denoise": self.extra_denoise,
            "pitch_change_all": self.pitch_change_all,
        }

    def update_activity(self):
        """Update last activity timestamp."""
        self.last_activity = time.time()


class StateManager:
    """
    Thread-safe user state manager.
    Manages per-user sessions and active processing tasks.
    """

    def __init__(self, session_timeout: int = 3600):
        """
        Args:
            session_timeout: Seconds before idle sessions are cleaned up (default 1 hour)
        """
        self._sessions: Dict[int, UserSession] = {}
        self._active_tasks: Dict[int, str] = {}  # user_id -> task_id
        self._lock = threading.Lock()
        self._session_timeout = session_timeout

    def get_session(self, user_id: int) -> UserSession:
        """Get or create a user session."""
        with self._lock:
            if user_id not in self._sessions:
                self._sessions[user_id] = UserSession()
            self._sessions[user_id].update_activity()
            return self._sessions[user_id]

    def set_state(self, user_id: int, state: UserState):
        """Set user's conversation state."""
        session = self.get_session(user_id)
        session.state = state
        session.update_activity()

    def get_state(self, user_id: int) -> UserState:
        """Get user's current state."""
        return self.get_session(user_id).state

    def clear_session(self, user_id: int):
        """Reset user session to defaults."""
        with self._lock:
            if user_id in self._sessions:
                self._sessions[user_id] = UserSession()

    def is_processing(self, user_id: int) -> bool:
        """Check if user has an active processing task."""
        return self.get_state(user_id) == UserState.PROCESSING

    def start_task(self, user_id: int, task_id: str):
        """Mark user as processing with a task ID."""
        with self._lock:
            self._active_tasks[user_id] = task_id
            if user_id in self._sessions:
                self._sessions[user_id].state = UserState.PROCESSING
                self._sessions[user_id].update_activity()

    def end_task(self, user_id: int) -> Optional[str]:
        """Mark user as done processing. Returns the task ID if any."""
        with self._lock:
            task_id = self._active_tasks.pop(user_id, None)
            if user_id in self._sessions:
                self._sessions[user_id].state = UserState.IDLE
                self._sessions[user_id].update_activity()
            return task_id

    def cleanup_stale_sessions(self):
        """Remove sessions that have been idle too long."""
        now = time.time()
        with self._lock:
            stale = [
                uid for uid, session in self._sessions.items()
                if now - session.last_activity > self._session_timeout
                and session.state == UserState.IDLE
            ]
            for uid in stale:
                del self._sessions[uid]
            if stale:
                return len(stale)
        return 0

    def get_active_count(self) -> int:
        """Count users currently processing."""
        with self._lock:
            return len(self._active_tasks)


# Global state manager instance
state_manager = StateManager()