from __future__ import annotations

import argparse
import subprocess
import tempfile
from pathlib import Path

import numpy as np
import soundfile as sf
from clearvoice import ClearVoice


MODEL = "MossFormer2_SE_48K"
SAMPLE_RATE = 48_000


def convert_input(source: Path, target: Path) -> None:
    completed = subprocess.run(
        [
            "ffmpeg", "-hide_banner", "-loglevel", "error", "-y",
            "-i", str(source), "-vn", "-ar", str(SAMPLE_RATE),
            "-c:a", "pcm_f32le", str(target),
        ],
        capture_output=True,
        text=True,
        timeout=1800,
    )
    if completed.returncode:
        raise RuntimeError(completed.stderr[-4000:] or "ffmpeg input conversion failed")


def enhance(source: Path, speech_path: Path, noise_path: Path) -> None:
    with tempfile.TemporaryDirectory(prefix="clearvoice-") as temp_dir:
        converted = Path(temp_dir) / "input-48k.wav"
        convert_input(source, converted)
        audio, sample_rate = sf.read(converted, dtype="float32", always_2d=True)
        if sample_rate != SAMPLE_RATE:
            raise RuntimeError(f"unexpected sample rate: {sample_rate}")

        model = ClearVoice(task="speech_enhancement", model_names=[MODEL])
        # Use ClearVoice's file-I/O path so recordings longer than its 20-second
        # one-pass window are segmented correctly. Run each channel separately
        # because the enhancement network itself is mono, then restore stereo.
        channels = []
        for channel_index in range(audio.shape[1]):
            channel_path = Path(temp_dir) / f"channel-{channel_index}.wav"
            sf.write(channel_path, audio[:, channel_index], SAMPLE_RATE, subtype="FLOAT")
            result = np.asarray(model(str(channel_path), False), dtype=np.float32).squeeze()
            if result.ndim != 1:
                raise RuntimeError(f"unexpected ClearVoice output shape: {result.shape}")
            channels.append(result)
        enhanced = np.column_stack(channels)

        length = min(len(audio), len(enhanced))
        original = audio[:length]
        speech = enhanced[:length]
        noise = original - speech

        # FLAC does not support floating-point samples. PCM_24 retains ample
        # headroom and avoids the invalid FLOAT/FLAC combination in libsndfile.
        sf.write(speech_path, np.clip(speech, -1.0, 1.0), SAMPLE_RATE, format="FLAC", subtype="PCM_24")
        sf.write(noise_path, np.clip(noise, -1.0, 1.0), SAMPLE_RATE, format="FLAC", subtype="PCM_24")


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("input", nargs="?", type=Path)
    parser.add_argument("speech", nargs="?", type=Path)
    parser.add_argument("noise", nargs="?", type=Path)
    parser.add_argument("--download-only", action="store_true")
    args = parser.parse_args()
    if args.download_only:
        ClearVoice(task="speech_enhancement", model_names=[MODEL])
        return
    if not all((args.input, args.speech, args.noise)):
        parser.error("input, speech and noise output paths are required")
    enhance(args.input, args.speech, args.noise)


if __name__ == "__main__":
    main()