"""Headless adapter for pinned deepbeepmeep/YuEGP inference functions.

No Comfy imports, no Gradio import/server, and no automatic profile fallback.
Stdout AIGEN_EVENT lines are consumed by the device-local host agent.
"""
import argparse
import ast
import hashlib
import gc
import json
import os
from pathlib import Path
import re
import sys
import time
import threading
from types import SimpleNamespace, FunctionType

YUEGP_REVISION = '2d72ff734b7a127324353c0dcd0f95ca4cc0b797'


def emit(**event):
    print('AIGEN_EVENT ' + json.dumps(event), flush=True)


def normalize_lyrics(text):
    # Accept our existing chips (Verse 1, Pre-Chorus); upstream only accepts \w+.
    sections = re.findall(r'\[([^\]]+)\]\s*([^\[]*)', text.strip(), re.S)
    sections = [(name, words.strip()) for name, words in sections if words.strip()]
    if len(sections) != 1:
        raise ValueError('YuEGP requires one non-empty lyric section. Combine the lyrics under one heading.')
    name, words = sections[0]
    return '[' + re.sub(r'\W+', '', name).lower() + ']\n' + words + '\n\n'


def load_functions(root):
    source = root / 'inference' / 'gradio_server.py'
    text = source.read_text(encoding='utf-8')
    if hashlib.sha256(text.encode()).hexdigest() != '567b91b3100c4fb2c95496e83b135e9d55fd03714e83912f9586037805fdecc9':
        raise RuntimeError('YuEGP inference source differs from the validated revision.')
    tree = ast.parse(text)
    names = {'BlockTokenRangeProcessor', 'load_audio_mono', 'encode_audio',
             'stage1_inference', 'stage2_generate', 'stage2_inference'}
    # Load the library functions only. Never execute upstream CLI/UI startup.
    nodes = [n for n in tree.body if isinstance(n, (ast.FunctionDef, ast.ClassDef)) and n.name in names]
    if {n.name for n in nodes} != names:
        raise RuntimeError('Unsupported YuEGP source: reinstall the pinned revision.')
    ns = {'__file__': str(source)}
    imports = [n for n in tree.body if isinstance(n, (ast.Import, ast.ImportFrom))
               and not (isinstance(n, ast.Import) and any(a.name == 'gradio' for a in n.names))]
    exec(compile(ast.Module(body=imports + nodes, type_ignores=[]), str(source), 'exec'), ns)
    return ns


def load_codec_on_cpu(ns, torch):
    with torch.device('cpu'):
        config = ns['OmegaConf'].load('xcodec_mini_infer/final_ckpt/config.yaml')
        codec = ns['SoundStream'](**config.generator.config)
        codec.load_state_dict(torch.load('xcodec_mini_infer/final_ckpt/ckpt_00360000.pth', map_location='cpu', weights_only=False)['codec_model'])
        codec.eval()
    return codec


def memory_profile_options(profile, compile_enabled, platform, total_ram):
    options = dict(profile_no=profile, quantizeTransformer=profile == 3,
                   compile=compile_enabled, verboseLevel=1)
    # Windows CUDA pinned allocations can fail before inference on 32GB hosts.
    # This affects transfer staging only, not weight precision or VRAM budgets.
    if platform == 'win32' and total_ram < 48 * 1024 ** 3:
        options['pinnedMemory'] = False
    return options


def main():
    import psutil
    parent = psutil.Process(os.getppid())
    def parent_watchdog():
        while parent.is_running():
            time.sleep(2)
        os._exit(2)  # The host died: never leave an orphan consuming the GPU.
    threading.Thread(target=parent_watchdog, daemon=True).start()
    parser = argparse.ArgumentParser()
    parser.add_argument('--root', required=True)
    parser.add_argument('--request', required=True)
    parser.add_argument('--profile', type=int, choices=[1, 3], default=1)
    parser.add_argument('--compile', action='store_true')
    cli = parser.parse_args()
    request = json.loads(Path(cli.request).read_text(encoding='utf-8-sig'))
    root = Path(cli.root).resolve()
    output = Path(cli.request).resolve().parent
    duration = int(request.get('duration', 60))
    if not 30 <= duration <= 150:
        raise ValueError('Duration must be 30–150 seconds.')
    lyrics = normalize_lyrics(request['lyrics'])
    tags = ' '.join(str(request['tags']).split())
    seed = int(request['seed'])
    max_tokens = duration * 100
    if max_tokens >= 16000:
        raise ValueError('YuEGP supports at most 150 seconds per section.')
    compile_enabled = False
    if cli.compile:
        import triton  # Explicit opt-in AND an actual successful import required.
        compile_enabled = True
    os.chdir(root / 'inference')
    sys.path[:0] = [str(root / 'inference'), str(root / 'inference/xcodec_mini_infer'),
                    str(root / 'inference/xcodec_mini_infer/descriptaudiocodec')]
    emit(stage='loading', progress=1, message=f'Loading YuEGP profile {cli.profile}', profile=cli.profile,
         compile=compile_enabled, sections=1, maxNewTokens=max_tokens)
    ns = load_functions(root)
    torch, np, sf = ns['torch'], ns['np'], ns['sf']
    if not torch.cuda.is_available():
        raise RuntimeError('YuEGP requires a CUDA GPU; CPU fallback is disabled.')
    # The upstream transformer patch is part of YuEGP, in its own environment.
    import transformers.generation.utils as generation
    if 'callback' not in Path(generation.__file__).read_text(encoding='utf-8'):
        raise RuntimeError('YuEGP transformer functions are missing. Run setup-yuegp.ps1.')
    attention = 'sdpa'
    try:
        import flash_attn
        attention = 'flash_attention_2'
    except ImportError:
        pass
    ns['random'].seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    device = torch.device('cuda:0')
    stage1 = request.get('stage1Model') or 'm-a-p/YuE-s1-7B-anneal-en-cot'
    stage2 = request.get('stage2Model') or 'm-a-p/YuE-s2-1B-general'
    model = ns['AutoModelForCausalLM'].from_pretrained(stage1, torch_dtype=torch.bfloat16, attn_implementation=attention).eval()
    model2 = ns['AutoModelForCausalLM'].from_pretrained(stage2, torch_dtype=torch.float16, attn_implementation=attention).eval()
    if not compile_enabled:
        # Transformers 4.48 selects DynamicCache with None; the literal
        # string "dynamic" is not an accepted configuration value.
        model.generation_config.cache_implementation = None
        model2.generation_config.cache_implementation = None
    model._validate_model_kwargs = lambda _: None
    model2._validate_model_kwargs = lambda _: None
    profile_options = memory_profile_options(cli.profile, compile_enabled, sys.platform, psutil.virtual_memory().total)
    emit(stage='loading', message='Preparing YuEGP memory', profileOptions=profile_options)
    offloader = ns['offload'].profile({'transformer': model, 'stage2': model2}, **profile_options)
    def memory_checkpoint(stage):
        torch.cuda.synchronize()
        free, total = torch.cuda.mem_get_info()
        emit(stage=stage, message=f'YuEGP {stage}', gpuFreeMiB=round(free / 1024 ** 2),
             gpuTotalMiB=round(total / 1024 ** 2), gpuAllocatedMiB=round(torch.cuda.memory_allocated() / 1024 ** 2),
             ramAvailableMiB=round(psutil.virtual_memory().available / 1024 ** 2))
    memory_checkpoint('memory-ready')
    args = SimpleNamespace(use_audio_prompt=False, use_dual_tracks_prompt=False, rescale=True,
                           output_dir=str(output), cuda_idx=0)
    ns.update(model=model, model_stage2=model2, device=device, codec_model=None, args=args,
              mmtokenizer=ns['_MMSentencePieceTokenizer']('./mm_tokenizer_v0.2_hf/tokenizer.model'),
              codectool=ns['CodecManipulator']('xcodec', 0, 1), codectool_stage2=ns['CodecManipulator']('xcodec', 0, 8),
              stage1_output_dir=str(output / 'stage1'),
              split_lyrics=lambda _: [lyrics], get_song_id=lambda *a: 'song')
    (output / 'stage1').mkdir(exist_ok=True)
    (output / 'stage2').mkdir(exist_ok=True)
    state = {}
    last = [0.0]
    def callback(done, total):
        now = time.monotonic()
        if now - last[0] < 1 and done < total:
            return
        last[0] = now
        stage = state.get('stage', 'Generating')
        # Percent is explicitly local to the reported stage, never a song ETA.
        emit(stage=stage, message=stage, step=int(done), maxStep=int(total),
             progress=round(100 * done / max(1, total), 1))
    emit(stage='stage1', message='Generating song tokens', progress=0)
    stems = ns['stage1_inference'](tags, lyrics, 1, max_tokens, seed, state, callback)
    memory_checkpoint('stage1-finished')
    emit(stage='stage2', message='Generating audio detail', progress=0)
    results = ns['stage2_inference'](model2, stems, str(output / 'stage2'),
                                    batch_size=20 if cli.profile == 1 else 4, state=state, callback=callback)
    # Release transformer allocations before codec/vocoder decode.
    offloader.unload_all()
    gc.collect()
    torch.cuda.empty_cache()
    # MMGP sets the default device to CUDA. Explicit CPU construction prevents
    # decoder initialization from competing with the language models for VRAM.
    emit(stage='decoding', message='Loading audio decoder', progress=0)
    codec = load_codec_on_cpu(ns, torch)
    codec.to(device)
    emit(stage='decoding', message='Decoding and mixing audio', progress=0)
    low_tracks = []
    for path in results:
        codes = np.load(path)
        with torch.no_grad():
            wave = codec.decode(torch.as_tensor(codes.astype(np.int16), dtype=torch.long).unsqueeze(0).permute(1, 0, 2).to(device))
        low_tracks.append(wave.cpu().squeeze().numpy())
    sf.write(str(output / 'mix16.wav'), low_tracks[0] + low_tracks[1], 16000)
    with torch.device('cpu'):
        vocal_decoder, inst_decoder = ns['build_codec_model']('xcodec_mini_infer/decoders/config.yaml',
            'xcodec_mini_infer/decoders/decoder_131000.pth', 'xcodec_mini_infer/decoders/decoder_151000.pth')
    # Keep upstream neural decoding, but write WAV directly: no MP3/FFmpeg
    # backend dependency in the isolated Windows environment.
    process_audio = ns['process_audio']
    audio_globals = dict(process_audio.__globals__)
    def save_wave(wave, path, sample_rate, rescale=False):
        data = wave.detach().cpu().numpy()
        peak = float(np.max(np.abs(data)))
        if rescale and peak > 0.99:
            data = data * (0.99 / peak)
        sf.write(str(path), data.T, sample_rate, subtype='PCM_16')
    audio_globals['save_audio'] = save_wave
    decode_audio = FunctionType(process_audio.__code__, audio_globals)
    tracks = []
    for path in results:
        instrumental = '_itrack' in path
        decoder = inst_decoder if instrumental else vocal_decoder
        with torch.no_grad():
            tracks.append(decode_audio(path, str(output / ('instrumental.wav' if instrumental else 'vocal.wav')),
                          True, args, decoder, codec))
        decoder.to('cpu')
        torch.cuda.empty_cache()
    mixed = (tracks[0] + tracks[1]).detach().cpu().squeeze().numpy()
    sf.write(str(output / 'mix44.wav'), mixed, 44100)
    ns['replace_low_freq_with_energy_matched'](a_file=str(output / 'mix16.wav'), b_file=str(output / 'mix44.wav'),
                                              c_file=str(output / 'audio.wav'), cutoff_freq=5500.0)
    info = sf.info(str(output / 'audio.wav'))
    if info.frames <= 0:
        raise RuntimeError('YuEGP produced empty audio.')
    emit(stage='complete', message='Audio ready', progress=100, duration=info.duration)


if __name__ == '__main__':
    try:
        main()
    except Exception as error:
        emit(stage='error', message=str(error), error=str(error))
        raise