"""Offline, bounded HTDemucs inference. Model installation is a separate operator step."""
import json
import os
import sys
from pathlib import Path
for key in ('OMP_NUM_THREADS', 'OPENBLAS_NUM_THREADS', 'MKL_NUM_THREADS'):
    os.environ[key] = '2'
CHECKSUM = "8726e21a993978c7ba086d3872e7608d7d5bfca646ca4aca459ffda844faa8b4"
SIGNATURE = '955717e8'

def available():
    import importlib.util
    root = Path(os.environ.get('SMART_EDIT_SEPARATION_MODEL_DIR', '/nonexistent'))
    return importlib.util.find_spec('demucs') is not None and any(root.glob(SIGNATURE + '-*.th'))

def separate(source, destination):
    if not available():
        raise RuntimeError('人声分离模型尚未配置，请安装 Demucs 依赖和官方 HTDemucs 权重。')
    import numpy as np
    import soundfile as sf
    import torch
    from demucs.states import load_model
    import hashlib
    from demucs.apply import apply_model
    torch.set_num_threads(2)
    torch.manual_seed(0)
    audio, sr = sf.read(source, dtype='float32', always_2d=True)
    if sr != 44100 or audio.shape[1] != 2 or not 0 < len(audio) <= sr * 180:
        raise ValueError('分离输入必须为 180 秒以内的 44.1kHz 双声道音频')
    if not np.isfinite(audio).all():
        raise ValueError('音频包含无效采样')
    checkpoint = next(Path(os.environ['SMART_EDIT_SEPARATION_MODEL_DIR']).glob(SIGNATURE + '-*.th'))
    if hashlib.sha256(checkpoint.read_bytes()).hexdigest() != CHECKSUM:
        raise ValueError('人声分离权重校验失败')
    # Official checkpoint verified above; Demucs 4 uses a pickled model class.
    model = load_model(torch.load(checkpoint, map_location='cpu', weights_only=False))
    model.eval()
    wave = torch.from_numpy(audio.T.copy())
    reference = wave.mean(0)
    mean, std = reference.mean(), reference.std()
    if std < 1e-8:
        vocals = np.zeros_like(audio)
    else:
        with torch.inference_mode():
            output = apply_model(model, ((wave - mean) / std)[None], device='cpu',
                                 shifts=1, split=True, overlap=.25, segment=7,
                                 progress=False, num_workers=0)[0]
            vocals = (output[model.sources.index('vocals')] * std + mean).T.numpy()
    # Complementary stems reconstruct the original exactly, with no independent gain normalization.
    destination = Path(destination)
    destination.mkdir(parents=True, exist_ok=True)
    sf.write(destination / 'vocals.wav', vocals, sr, subtype='FLOAT')
    sf.write(destination / 'background.wav', audio - vocals, sr, subtype='FLOAT')

if __name__ == '__main__':
    if sys.argv[1:] == ['--check']:
        print(json.dumps({'separation': available()}))
    else:
        separate(*sys.argv[1:3])