import src.utils.config_manager as config

config = config.config

MODEL_NAME = "model_name"
MODEL_PATH = "model_path"
DEVICE = "device"
DTYPE = "dtype"
NUM_STEPS = "num_steps"
SPEED = "speed"
PROFILES_DIR = "profiles_dir"
OUTPUT_DIR = "output_dir"
ENABLED = "enabled"

DEFAULT_MODEL_NAME = "k2-fsa/OmniVoice"
DEFAULT_MODEL_PATH = "models/OmniVoice"
DEFAULT_DEVICE = "cuda:0"
DEFAULT_DTYPE = "float16"
DEFAULT_NUM_STEPS = 32
DEFAULT_SPEED = 1.0
DEFAULT_PROFILES_DIR = "data/voice_profiles"
DEFAULT_OUTPUT_DIR = "output/voice"
DEFAULT_ENABLED = True


def register_params():
    config.register(
        name=MODEL_NAME,
        val=DEFAULT_MODEL_NAME,
        cat="omnivoice",
        desc="Model name (HuggingFace) or use model_path for local",
    )
    config.register(
        name=MODEL_PATH,
        val=DEFAULT_MODEL_PATH,
        cat="omnivoice",
        desc="Local model path (if set, overrides model_name)",
    )
    config.register(
        name=DEVICE, val=DEFAULT_DEVICE, cat="omnivoice", desc="Device for inference"
    )
    config.register(name=DTYPE, val=DEFAULT_DTYPE, cat="omnivoice", desc="Data type")
    config.register(
        name=NUM_STEPS, val=DEFAULT_NUM_STEPS, cat="omnivoice", desc="Diffusion steps"
    )
    config.register(name=SPEED, val=DEFAULT_SPEED, cat="omnivoice", desc="Speed factor")
    config.register(
        name=PROFILES_DIR,
        val=DEFAULT_PROFILES_DIR,
        cat="omnivoice",
        desc="Voice profiles directory",
    )
    config.register(
        name=OUTPUT_DIR,
        val=DEFAULT_OUTPUT_DIR,
        cat="omnivoice",
        desc="Output directory for audio",
    )
    config.register(
        name=ENABLED,
        val=DEFAULT_ENABLED,
        cat="omnivoice",
        desc="Enable omnivoice module",
    )