"""Humaneness Voice Large 1.3 Alpha: portable full-model inference.

Use a downloaded model directory, or a Hugging Face model ID. This file has no
JUPITER-specific runtime dependency. Twelve MOSS-v2 codebooks,48kHz mono.
"""
import argparse
import json
from pathlib import Path

LANGUAGES = {'en': 'English', 'de': 'German', 'nl': 'Dutch', 'fr': 'French', 'es': 'Spanish'}


def load_model(directory, device='cuda:0'):
    import torch
    from transformers import AutoModel, AutoTokenizer
    from transformers.dynamic_module_utils import get_class_from_dynamic_module
    model = AutoModel.from_pretrained(str(directory), trust_remote_code=True,
        dtype=torch.bfloat16, attn_implementation='sdpa').to(device).eval()
    Processor = get_class_from_dynamic_module('processing_moss_tts.MossTTSLocalProcessor', str(directory))
    tokenizer = AutoTokenizer.from_pretrained(str(directory), trust_remote_code=True)
    proc = Processor(tokenizer=tokenizer, audio_tokenizer=None, model_config=model.config)
    return model, proc


def load_codec(proc, path, device='cuda:0'):
    from transformers import AutoModel
    codec = AutoModel.from_pretrained(str(path), trust_remote_code=True, codec_weight_dtype='fp32')
    if hasattr(codec, 'set_compute_dtype'):
        codec.set_compute_dtype('bf16')
    proc.audio_tokenizer = codec.to(device).eval()
    return proc


def build_user(proc, *, text, prompt, frames, language, reference=None, alias=None):
    if not text.strip() or not prompt.strip() or language not in LANGUAGES:
        raise ValueError('Nonempty literal text, English directions and a supported language required')
    if reference is not None and (reference.ndim != 2 or reference.shape[1] != 12
            or not 63 <= len(reference) <= 187):
        raise ValueError('Reference codes must contain12codebooks over5.04–14.96seconds')
    if alias is not None:
        if reference is not None:
            raise ValueError('Choose either a reference recording or a registered voice name, not both')
        catalog = json.loads(Path(__file__).with_name('voices.json').read_text())
        names = {v['speaker_alias'] for v in catalog['voices']}
        if alias not in names:
            raise ValueError('Unknown voice name; see VOICE_ALIASES.md and voices.json')
    user = proc.build_user_message(text=text, instruction=prompt, tokens=int(frames),
        language=LANGUAGES[language], reference=[reference] if reference is not None else None)
    if alias is not None:
        # Exact alias representation used by training/legacy/acting_data.py.
        # It is not an invented prefix in the spoken text or a new special token.
        marker = '- Reference(s):\nNone\n'
        if user['content'].count(marker) != 1 or user.get('audio_codes_list'):
            raise RuntimeError('Processor wrapper no longer matches the trained voice-name format')
        user['content'] = user['content'].replace(marker, f'- Reference(s):\nSpeaker: {alias}\n', 1)
    return user


def generate(model, proc, user, *, seed=777, max_frames=625, device='cuda:0'):
    import torch
    batch = proc([[user]], mode='generation')
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    with torch.inference_mode(), torch.autocast('cuda', dtype=torch.bfloat16):
        output = model.generate(input_ids=batch['input_ids'].to(device),
            attention_mask=batch['attention_mask'].to(device), max_new_frames=max_frames,
            do_sample=True, audio_temperature=1., audio_top_p=.95, audio_top_k=50,
            audio_repetition_penalty=1., use_kv_cache=True)
    prefix, sequence = output[0]
    if prefix != 0 or sequence.ndim != 2 or sequence.shape[1] != 13:
        raise RuntimeError('Unexpected MOSS frame layout')
    codes = sequence[1:, 1:]
    codes = codes[(codes[:, 0] >= 0) & (codes[:, 0] != model.config.audio_pad_token_id)]
    if len(codes) < 2 or int(codes.min()) < 0 or int(codes.max()) > 1023:
        raise RuntimeError('Too few frames or corrupt codec indices')
    return codes.cpu()


def decode(proc, codes, device='cuda:0'):
    import numpy as np
    import torch
    with torch.inference_mode():
        wave = proc.decode_audio_codes([codes.to(device)], return_stereo=False)[0]
    wave = wave.float().cpu().numpy().reshape(-1)
    rate = int(proc.model_config.sampling_rate)
    if not np.isfinite(wave).all() or abs(len(wave) / len(codes) - rate / 12.5) > .01 * rate / 12.5:
        raise RuntimeError('Codec frame-rate/sample-rate canary failed')
    return wave, rate


def main():
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument('--model', default='laion/Humaneness-Voice-Large-1.3-Alpha')
    p.add_argument('--codec', default='OpenMOSS-Team/MOSS-Audio-Tokenizer-v2')
    p.add_argument('--text', required=True, help='Literal spoken words, without directing tags')
    p.add_argument('--prompt', help='GENERAL/SCRIPT directions; literal \\n accepted')
    p.add_argument('--prompt-file', type=Path)
    p.add_argument('--language', choices=LANGUAGES, default='en')
    voice = p.add_mutually_exclusive_group()
    voice.add_argument('--reference-wav', type=Path)
    voice.add_argument('--voice-name', help='Registered fictional voice alias; see VOICE_ALIASES.md')
    p.add_argument('--seconds', type=float, default=8.)
    p.add_argument('--seed', type=int, default=777)
    p.add_argument('--output', type=Path, required=True)
    args = p.parse_args()
    if bool(args.prompt) == bool(args.prompt_file):
        p.error('Supply exactly one of --prompt or --prompt-file')
    if not 1 <= args.seconds <= 60:
        p.error('Target seconds must be between1and60')
    import numpy as np
    import soundfile as sf
    import torch
    if not torch.cuda.is_available():
        raise RuntimeError('The tested path requires a CUDA GPU')
    device = 'cuda:0'
    model, proc = load_model(args.model, device)
    load_codec(proc, args.codec, device)
    reference = None
    if args.reference_wav:
        wave, sr = sf.read(args.reference_wav, dtype='float32', always_2d=True)
        if not np.isfinite(wave).all() or len(wave) / sr < 5.04:
            raise ValueError('Use at least5.04seconds of clean single-speaker reference speech')
        # Never reinstate37frames/2.96s. Crop waveform BEFORE encoding.
        wave = wave[:int(sr * 14.96)]
        with torch.inference_mode():
            reference = proc.encode_audios_from_wav([torch.from_numpy(wave.T.copy())], int(sr), n_vq=12)[0][:187]
    prompt = args.prompt_file.read_text() if args.prompt_file else args.prompt.replace('\\n', '\n')
    frames = round(args.seconds * 12.5)
    user = build_user(proc, text=args.text, prompt=prompt, frames=frames,
        language=args.language, reference=reference, alias=args.voice_name)
    codes = generate(model, proc, user, seed=args.seed, max_frames=min(1000, frames + 190))
    wave, sr = decode(proc, codes)
    args.output.parent.mkdir(parents=True, exist_ok=True)
    if args.output.exists():
        raise FileExistsError('Refusing to overwrite ' + str(args.output))
    sf.write(args.output, wave, sr, subtype='PCM_16')
    print(json.dumps({'output': str(args.output), 'duration_s': len(wave) / sr,
        'requested_s': args.seconds, 'frames': len(codes), 'language': args.language,
        'reference_frames': len(reference) if reference is not None else 0,
        'voice_name': args.voice_name, 'seed': args.seed}))


if __name__ == '__main__':
    main()