from __future__ import annotations

import base64
from contextlib import redirect_stdout
import io
import json
import math
from pathlib import Path
import sys


def load_reference_audio(path):
    # torchaudio.load on the base requires TorchCodec. Decode this local input
    # with libsndfile instead; torchaudio functional resampling remains intact.
    import soundfile as sf
    import torch
    info = sf.info(str(path))
    if not 0 < info.duration <= 120 or not 1 <= info.channels <= 2 or not 8000 <= info.samplerate <= 192000:
        raise ValueError('Reference audio must be 0-120 seconds, mono/stereo, 8-192 kHz')
    audio, rate = sf.read(str(path), dtype='float32', always_2d=True)
    if not len(audio) or not rate or not torch.isfinite(torch.from_numpy(audio)).all():
        raise ValueError('Reference audio is empty or invalid')
    return torch.from_numpy(audio.T.copy()), int(rate)


def synthesize(data):
    import soundfile as sf
    import torch
    sys.path.insert(0, '/opt/irodori-tts')
    from irodori_tts import inference_runtime as runtime_module
    runtime_module._load_audio = load_reference_audio
    runtime_module._resolve_tokenizer_source = lambda checkpoint, fallback: (data['tokenizer'], True)
    checkpoint = Path(data['checkpoint'])
    if not checkpoint.is_file():
        raise ValueError('Irodori checkpoint does not exist')
    reference = data.get('reference_wav')
    references = data.get('reference_wavs')
    if reference and references:
        raise ValueError('Use one reference representation')
    paths = references if references is not None else ([reference] if reference else [])
    if not isinstance(paths, list) or len(paths) > 20 or any(not isinstance(path, str) or not path for path in paths):
        raise ValueError('Invalid reference audio paths')
    total_seconds = 0
    for path in paths:
        info = sf.info(path)
        if not 0 < info.duration <= 120 or not 1 <= info.channels <= 2 or not 8000 <= info.samplerate <= 192000:
            raise ValueError('Reference audio must be 0-120 seconds, mono/stereo, 8-192 kHz')
        total_seconds += info.duration
    if total_seconds > 120:
        raise ValueError('Reference audio must be at most 120 seconds combined')
    decoded = {str(path): load_reference_audio(path) for path in paths}
    if decoded:
        runtime_module._load_audio = lambda path: decoded[str(path)] if str(path) in decoded else load_reference_audio(path)
    has_reference = bool(paths)
    caption = data.get('caption') or None
    runtime, reloaded = runtime_module.get_cached_runtime(runtime_module.RuntimeKey(
        checkpoint=str(checkpoint), model_device='cuda', model_precision='bf16',
        codec_repo=data['codec'], codec_device='cuda',
        codec_precision='fp32', codec_deterministic_encode=True,
        codec_deterministic_decode=True, compile_model=False, compile_dynamic=False,
    ))
    if caption and not runtime.model_cfg.use_caption_condition:
        raise ValueError('Selected model does not support voice descriptions')
    if has_reference and not runtime.model_cfg.use_speaker_condition_resolved:
        raise ValueError('Selected model does not support reference audio')
    text_cfg, caption_cfg, speaker_cfg, messages = runtime_module.resolve_cfg_scales(
        cfg_guidance_mode='independent', cfg_scale_text=data['cfg_scale_text'],
        cfg_scale_caption=data['cfg_scale_caption'] if caption else 0,
        cfg_scale_speaker=data['cfg_scale_speaker'] if has_reference else 0,
        cfg_scale=None, use_caption_condition=bool(caption), use_speaker_condition=has_reference,
    )
    request = runtime_module.SamplingRequest(
        text=data['text'], caption=caption, ref_wav=reference, ref_wavs=references, no_ref=not has_reference,
        seconds=data.get('seconds'), duration_scale=data['duration_scale'],
        min_seconds=0.5, max_seconds=45, max_ref_seconds=120,
        num_candidates=1, decode_mode='sequential', num_steps=data['num_steps'],
        seed=data['seed'], cfg_guidance_mode='independent', cfg_scale_text=text_cfg,
        cfg_scale_caption=caption_cfg, cfg_scale_speaker=speaker_cfg,
        context_kv_cache=True, trim_tail=True,
    )
    result = runtime.synthesize(request)
    audio = result.audio.detach().float().cpu().squeeze()
    if audio.ndim != 1 or not audio.numel() or not torch.isfinite(audio).all():
        raise ValueError('Irodori generated invalid audio')
    buffer = io.BytesIO()
    sf.write(buffer, audio.numpy(), int(result.sample_rate), format='WAV', subtype='PCM_16')
    Path(data['output']).write_bytes(buffer.getvalue())
    return {
        'backend': 'irodori',
        'audio_mime': 'audio/wav', 'sample_rate': int(result.sample_rate),
        'duration_seconds': audio.numel() / int(result.sample_rate),
        'seed': int(result.used_seed), 'num_steps': data['num_steps'],
        'model': checkpoint.name, 'runtime_reloaded': bool(reloaded),
        'messages': [*messages, *result.messages],
    }


if __name__ == '__main__':
    try:
        with redirect_stdout(sys.stderr):
            result = synthesize(json.load(sys.stdin))
    except Exception as exc:
        result = {'error': str(exc)}
    json.dump(result, sys.stdout, ensure_ascii=False)