src/config.py

"""
Configuration management module for the AIGC Telegram Bot.
Loads settings from environment variables and .env file.
"""

import os
from dataclasses import dataclass, field
from typing import List, Optional
from pathlib import Path

try:
    from dotenv import load_dotenv
    load_dotenv()
except ImportError:
    pass


@dataclass
class PipelineConfig:
    """Configuration for the AI Cover Generation pipeline."""
    max_concurrent_jobs: int = 2
    max_file_size_mb: int = 20
    default_output_format: str = "mp3"

    # RVC defaults
    default_pitch_change: int = 0
    default_index_rate: float = 0.5
    default_filter_radius: int = 3
    default_volume_envelope: float = 0.25
    default_f0_method: str = "rmvpe"
    default_hop_length: int = 128
    default_protect: float = 0.33

    # Audio effects defaults
    default_reverb_size: float = 0.15
    default_reverb_wet: float = 0.2
    default_reverb_dry: float = 0.8
    default_reverb_damping: float = 0.7

    # Mix gains
    default_main_gain: int = 0
    default_backup_gain: int = 0
    default_inst_gain: int = 0

    # Inference mode
    default_inference_mode: str = "full"


@dataclass
class BotConfig:
    """Configuration for the Telegram Bot."""
    bot_token: str = ""
    admin_ids: List[int] = field(default_factory=list)
    allowed_user_ids: Optional[List[int]] = None  # None = public bot
    bot_name: str = "AIGC Cover Bot"

    # Server
    webhook_host: str = "0.0.0.0"
    webhook_port: int = 8080
    webhook_url: Optional[str] = None

    # Logging
    log_level: str = "INFO"
    log_file: Optional[str] = None


@dataclass
class PathConfig:
    """Filesystem paths used by the bot."""
    base_dir: Path = Path(__file__).parent.parent
    models_dir: Path = Path(__file__).parent.parent / "data" / "models"
    temp_dir: Path = Path(__file__).parent.parent / "data" / "temp"
    output_dir: Path = Path(__file__).parent.parent / "data" / "output"
    mdx_models_dir: Path = Path(__file__).parent.parent / "assets" / "mdxnet_models"

    def __post_init__(self):
        """Ensure all required directories exist."""
        for dir_path in [self.models_dir, self.temp_dir, self.output_dir, self.mdx_models_dir]:
            dir_path.mkdir(parents=True, exist_ok=True)


def _parse_int_list(value: str) -> List[int]:
    """Parse a comma-separated string of integers."""
    if not value or not value.strip():
        return []
    return [int(x.strip()) for x in value.split(",") if x.strip().isdigit()]


def load_config() -> BotConfig:
    """Load bot configuration from environment variables."""
    allowed_raw = os.getenv("ALLOWED_USER_IDS", "")
    allowed_ids = _parse_int_list(allowed_raw) if allowed_raw.strip() else None

    return BotConfig(
        bot_token=os.getenv("BOT_TOKEN", ""),
        admin_ids=_parse_int_list(os.getenv("ADMIN_IDS", "")),
        allowed_user_ids=allowed_ids,
        bot_name=os.getenv("BOT_NAME", "AIGC Cover Bot"),
        webhook_host=os.getenv("WEBHOOK_HOST", "0.0.0.0"),
        webhook_port=int(os.getenv("WEBHOOK_PORT", "8080")),
        webhook_url=os.getenv("WEBHOOK_URL", "") or None,
        log_level=os.getenv("LOG_LEVEL", "INFO"),
        log_file=os.getenv("LOG_FILE", "") or None,
    )


def load_pipeline_config() -> PipelineConfig:
    """Load pipeline configuration from environment variables."""
    return PipelineConfig(
        max_concurrent_jobs=int(os.getenv("MAX_CONCURRENT_JOBS", "2")),
        max_file_size_mb=int(os.getenv("MAX_FILE_SIZE_MB", "20")),
        default_output_format=os.getenv("DEFAULT_OUTPUT_FORMAT", "mp3"),
        default_pitch_change=int(os.getenv("DEFAULT_PITCH_CHANGE", "0")),
        default_index_rate=float(os.getenv("DEFAULT_INDEX_RATE", "0.5")),
        default_filter_radius=int(os.getenv("DEFAULT_FILTER_RADIUS", "3")),
        default_volume_envelope=float(os.getenv("DEFAULT_VOLUME_ENVELOPE", "0.25")),
        default_f0_method=os.getenv("DEFAULT_F0_METHOD", "rmvpe"),
        default_hop_length=int(os.getenv("DEFAULT_HOP_LENGTH", "128")),
        default_protect=float(os.getenv("DEFAULT_PROTECT", "0.33")),
        default_reverb_size=float(os.getenv("DEFAULT_REVERB_SIZE", "0.15")),
        default_reverb_wet=float(os.getenv("DEFAULT_REVERB_WET", "0.2")),
        default_reverb_dry=float(os.getenv("DEFAULT_REVERB_DRY", "0.8")),
        default_reverb_damping=float(os.getenv("DEFAULT_REVERB_DAMPING", "0.7")),
        default_main_gain=int(os.getenv("DEFAULT_MAIN_GAIN", "0")),
        default_backup_gain=int(os.getenv("DEFAULT_BACKUP_GAIN", "0")),
        default_inst_gain=int(os.getenv("DEFAULT_INST_GAIN", "0")),
        default_inference_mode=os.getenv("DEFAULT_INFERENCE_MODE", "full"),
    )


def get_path_config() -> PathConfig:
    """Get path configuration."""
    return PathConfig()