import torch
import torchaudio
from transformers import WhisperProcessor, WhisperForConditionalGeneration


_processor = None
_model = None


def _load_whisper():
    global _processor, _model
    if _model is None:
        _processor = WhisperProcessor.from_pretrained("openai/whisper-large-v3")
        _model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-large-v3")
    return _processor, _model


def isolate_vocals(path):
    """
    run demucs to split stems, return path to vocals.
    expects demucs CLI installed via pip.
    """
    import subprocess
    import os

    out_dir = os.path.join(os.path.dirname(path), '_stems')
    cmd = ['python', '-m', 'demucs', '--two-stems', 'vocals',
           '-o', out_dir, path]
    subprocess.run(cmd, check=True, capture_output=True)

    # demucs outputs to out_dir/htdemucs/<trackname>/vocals.wav
    track_name = os.path.splitext(os.path.basename(path))[0]
    vocals_path = os.path.join(out_dir, 'htdemucs', track_name, 'vocals.wav')

    if not os.path.exists(vocals_path):
        raise FileNotFoundError(f"demucs didn't produce vocals at {vocals_path}")

    return vocals_path


def transcribe(path, isolate=True):
    """
    extract lyrics from audio.
    if isolate=True, runs demucs first to pull vocals.
    """
    if isolate:
        try:
            vocal_path = isolate_vocals(path)
        except Exception:
            # fall back to raw audio if stem separation fails
            vocal_path = path
    else:
        vocal_path = path

    processor, model = _load_whisper()

    track, sr = torchaudio.load(vocal_path)

    # whisper wants 16kHz mono
    if sr != 16000:
        track = torchaudio.transforms.Resample(sr, 16000)(track)
    if track.shape[0] > 1:
        track = track.mean(dim=0, keepdim=True)

    track = track.squeeze()

    inputs = processor(track.numpy(), sampling_rate=16000, return_tensors="pt")

    with torch.no_grad():
        predicted_ids = model.generate(inputs.input_features)

    lyrics = processor.batch_decode(predicted_ids, skip_special_tokens=True)[0]
    return lyrics.strip()