main.py

"""
MellowMic - Real-time voice styles for creators
GUI mode (default) or CLI mode with --cli flag.

GUI:  python main.py
CLI:  python main.py --cli [--engine dsp] [--effect feminine_dsp] [--list-devices] ...
"""
import sys
import os

from core.branding import APP_NAME, APP_VERSION


def _run_gui() -> None:
    from PyQt5.QtWidgets import QApplication
    from PyQt5.QtGui import QFont
    from ui.main_window import MainWindow

    app = QApplication(sys.argv)
    app.setApplicationName(APP_NAME)
    app.setApplicationVersion(APP_VERSION)

    # Default font
    font = QFont("Segoe UI", 10)
    app.setFont(font)

    window = MainWindow()
    window.show()

    sys.exit(app.exec_())


def _run_cli() -> None:
    import argparse
    import signal
    import time

    import numpy as np
    import yaml

    from core.effects import PRESET_KEYS

    parser = argparse.ArgumentParser(description=f"{APP_NAME} CLI")
    parser.add_argument("--cli", action="store_true")
    parser.add_argument("--engine", choices=["dsp", "ai", "bypass"], default=None,
                        help="Processing engine")
    parser.add_argument("--effect", choices=PRESET_KEYS, default="bypass",
                        help="DSP voice effect preset")
    parser.add_argument("--model-id", type=str, default=None,
                        help="AI model bundle id under models/")
    parser.add_argument("--ai-provider", choices=["CUDAExecutionProvider", "CPUExecutionProvider"],
                        default=None, help="ONNX Runtime provider for AI mode")
    parser.add_argument("--chunk-ms", type=int, default=None,
                        help="AI chunk target in milliseconds")
    parser.add_argument("--crossfade-ms", type=int, default=None,
                        help="AI crossfade in milliseconds")
    parser.add_argument("--list-models", action="store_true")
    parser.add_argument("--list-devices", action="store_true")
    parser.add_argument("--input-device", type=str, default=None)
    parser.add_argument("--output-device", type=str, default=None)
    parser.add_argument("--no-vad", action="store_true")
    parser.add_argument("--config", type=str, default="config.yaml")
    args = parser.parse_args()

    try:
        with open(args.config) as f:
            config = yaml.safe_load(f)
    except FileNotFoundError:
        config = {}

    audio_cfg  = config.get("audio", {})
    sr         = audio_cfg.get("sample_rate", 48000)
    block_size = audio_cfg.get("block_size", 256)
    ring_depth = audio_cfg.get("ring_buffer_depth", 4)
    defaults   = config.get("defaults", {})
    ai_cfg     = config.get("ai", {})
    effects_cfg = config.get("effects", {})

    # Activate formant-preserving pitch (Pedalboard) if enabled in config.
    if effects_cfg.get("formant_preserving_pitch", False):
        from core.effects import enable_native_pitch
        enable_native_pitch()

    engine_mode = args.engine or defaults.get("engine", "dsp")
    ai_provider = args.ai_provider or ai_cfg.get("provider", "CUDAExecutionProvider")
    chunk_ms    = args.chunk_ms or ai_cfg.get("chunk_ms", 180)
    crossfade_ms = args.crossfade_ms or ai_cfg.get("crossfade_ms", 12)
    model_id    = args.model_id or ai_cfg.get("model_id")

    if args.list_devices:
        from core.routing import list_devices
        list_devices()
        sys.exit(0)

    from core.model_registry import ModelRegistry
    registry = ModelRegistry(config.get("models", {}).get("root", "models"))

    if args.list_models:
        try:
            models = registry.refresh()
        except Exception as e:
            print(f"[ERROR] Model registry: {e}")
            sys.exit(2)
        if not models:
            print("No licensed model bundles found in models/")
            sys.exit(0)
        for mid, meta in models.items():
            print(f"{mid}: {meta.name} ({meta.engine}, {meta.sample_rate} Hz)")
        sys.exit(0)

    input_dev = None
    if args.input_device:
        from core.routing import find_device
        input_dev = find_device(args.input_device, kind="input")

    output_dev = None
    if args.output_device:
        from core.routing import find_device
        output_dev = find_device(args.output_device, kind="output")
    if output_dev is None:
        from core.routing import find_vb_cable
        output_dev = find_vb_cable()
        if output_dev is not None:
            import sounddevice as sd
            print(f"[INFO] Auto-selected VB-Cable: [{output_dev}] {sd.query_devices(output_dev)['name']}")

    from core.vad import make_vad

    vad_cfg = config.get("vad", {})
    vad     = make_vad(
        backend=vad_cfg.get("backend", "webrtc"),
        aggressiveness=vad_cfg.get("aggressiveness", 2),
        frame_ms=vad_cfg.get("frame_ms", 20),
        pad_silence_ms=vad_cfg.get("pad_silence_ms", 200),
        sample_rate=sr,
        block_size=block_size,
        silero_model_path=vad_cfg.get("silero_model_path"),
        silero_threshold=vad_cfg.get("silero_threshold", 0.5),
    ) if not args.no_vad else None

    from core.engine_factory import create_voice_engine
    from core.noise_suppression import make_noise_suppressor
    from core.processor_chain import VoiceProcessorChain
    from core.voice_engine import VoiceEngineParams, VoiceEngineUnavailable

    params = VoiceEngineParams(
        preset=args.effect,
        model_id=model_id,
        provider=ai_provider,
        chunk_ms=chunk_ms,
        crossfade_ms=crossfade_ms,
    )
    try:
        voice_engine = create_voice_engine(
            engine_mode,
            params,
            sample_rate=sr,
            model_registry=registry,
        )
    except VoiceEngineUnavailable as e:
        print(f"[ERROR] {e}")
        sys.exit(2)

    noise = make_noise_suppressor(sample_rate=sr, prop_decrease=0.55)
    chain = VoiceProcessorChain(vad=vad, noise_suppressor=noise, voice_engine=voice_engine)

    def processor(block: np.ndarray) -> np.ndarray:
        return chain.process(block)

    block_ms = block_size / sr * 1000.0
    print(f"\n  {APP_NAME} - Voice Effects CLI")
    print(f"  Engine  : {engine_mode}")
    print(f"  Effect  : {args.effect if engine_mode == 'dsp' else '-'}")
    print(f"  Model   : {model_id or '-'}")
    print(f"  Provider: {voice_engine.stats().provider or ai_provider}")
    print(f"  Block   : {block_size} samples ({block_ms:.1f}ms @ {sr}Hz)")
    print(f"  VAD     : {'enabled' if not args.no_vad else 'disabled'}")
    print(f"\n  Press Ctrl+C to stop.\n")

    from core.audio_io import AudioEngine

    engine = AudioEngine(
        processor=processor,
        sample_rate=sr,
        block_size=block_size,
        input_device=input_dev,
        output_device=output_dev,
        ring_buffer_depth=ring_depth,
    )

    def _shutdown(sig, frame):
        print(f"\n  Stopping {APP_NAME}...")
        engine.stop()
        sys.exit(0)

    signal.signal(signal.SIGINT, _shutdown)
    engine.start()
    print("  [RUNNING] Speak into your mic...")

    while True:
        time.sleep(1)


if __name__ == "__main__":
    if "--cli" in sys.argv:
        _run_cli()
    else:
        _run_gui()