#!/usr/bin/env python3
"""Make one spoken line for the MetaHuman lab.

  ElevenLabs text-to-speech (--elevenlabs <voice_id>, key in ELEVENLABS_API_KEY), or --audio <file> you
  generated elsewhere, or — placeholder only — macOS `say`
    -> assets/speech/<id>.m4a          web audio
  PocketSphinx forced alignment of the transcript -> phones -> Rhubarb-style mouth cues
    -> assets/speech/<id>.json         {"metadata": {...}, "mouthCues": [{start, end, value}]}
  and registers the line in assets/speech/lines.json (loaded by Game.ts; press T in the game).

Usage (from the project root, with the venv in tools/lipsync/.venv):
  tools/lipsync/.venv/bin/python tools/lipsync/make_line.py intro "Hi, I'm Bo." [--voice Samantha]
  tools/lipsync/.venv/bin/python tools/lipsync/make_line.py greet "<transcript>" --audio path/to/tts.mp3
  ELEVENLABS_API_KEY=... tools/lipsync/.venv/bin/python tools/lipsync/make_line.py pitch "<text>" --elevenlabs <voice_id> [--model eleven_multilingual_v2]
"""
import argparse, base64, json, os, re, subprocess, sys, urllib.request, wave

ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
OUT = os.path.join(ROOT, 'assets', 'speech')

# ARPAbet phone -> Rhubarb mouth shape (A closed MBP, B teeth/most consonants, C open EH,
# D wide AA, E rounded AO/ER, F puckered UW/OW/W, G F/V, H L, X rest).
PHONE_SHAPE = {
    'AA': 'D', 'AE': 'C', 'AH': 'C', 'AO': 'E', 'AW': 'D', 'AY': 'D', 'EH': 'C', 'ER': 'E', 'EY': 'C',
    'IH': 'B', 'IY': 'B', 'OW': 'E', 'OY': 'E', 'UH': 'F', 'UW': 'F',
    'B': 'A', 'M': 'A', 'P': 'A', 'F': 'G', 'V': 'G', 'L': 'H', 'W': 'F', 'R': 'E',
    'SIL': 'X', '<sil>': 'X', '<s>': 'X', '</s>': 'X', '+SPN+': 'X', '+NSN+': 'X', '+BREATH+': 'X',
}
DEFAULT_SHAPE = 'B'
MIN_CUE = 0.03  # seconds; shorter cues are merged into their predecessor

# Pronunciations for words the CMU dictionary lacks. Extend as your lines need it.
EXTRA_WORDS = {
    'metahuman': 'M EH T AH HH Y UW M AH N', 'metahumans': 'M EH T AH HH Y UW M AH N Z',
    'bitmagic': 'B IH T M AE JH IH K', 'arkit': 'AA R K IH T',
    'blendshape': 'B L EH N D SH EY P', 'blendshapes': 'B L EH N D SH EY P S',
    'voxel': 'V AA K S AH L', 'voxels': 'V AA K S AH L Z',
}
ONES = ['zero', 'one', 'two', 'three', 'four', 'five', 'six', 'seven', 'eight', 'nine', 'ten', 'eleven', 'twelve',
        'thirteen', 'fourteen', 'fifteen', 'sixteen', 'seventeen', 'eighteen', 'nineteen']
TENS = ['', '', 'twenty', 'thirty', 'forty', 'fifty', 'sixty', 'seventy', 'eighty', 'ninety']

def num_words(n):
    if n < 20: return ONES[n]
    if n < 100: return TENS[n // 10] + ('' if n % 10 == 0 else ' ' + ONES[n % 10])
    if n < 1000: return ONES[n // 100] + ' hundred' + ('' if n % 100 == 0 else ' ' + num_words(n % 100))
    return ' '.join(ONES[int(c)] for c in str(n))

def normalize(text):
    t = text.lower().replace('-', ' ').replace('—', ' ').replace('…', ' ')
    t = re.sub(r'\d+', lambda m: ' ' + num_words(int(m.group())) + ' ', t)
    t = re.sub(r"[^a-z' ]+", ' ', t)
    return [w for w in t.split() if w]

def run(cmd, **kw):
    subprocess.run(cmd, check=True, **kw)

def elevenlabs_tts(args, base):
    """ElevenLabs `with-timestamps`: the audio plus character-level alignment (kept next to the line for
    tooling; the mouth cues still come from the forced alignment below, which yields phones)."""
    key = os.environ.get('ELEVENLABS_API_KEY')
    if not key: sys.exit('[make_line] --elevenlabs needs ELEVENLABS_API_KEY in the environment')
    api = (os.environ.get('ELEVENLABS_API_URL') or 'https://api.elevenlabs.io').rstrip('/')
    body = json.dumps({'text': args.text, 'model_id': args.model, 'output_format': 'mp3_44100_128'}).encode()
    req = urllib.request.Request(f'{api}/v1/text-to-speech/{args.elevenlabs}/with-timestamps', data=body,
                                 headers={'xi-api-key': key, 'Content-Type': 'application/json', 'Accept': 'application/json'})
    with urllib.request.urlopen(req, timeout=120) as res:
        payload = json.load(res)
    mp3 = base + '.source.mp3'
    with open(mp3, 'wb') as f: f.write(base64.b64decode(payload['audio_base64']))
    with open(base + '.alignment.json', 'w') as f: json.dump(payload.get('alignment') or payload.get('normalized_alignment') or {}, f)
    print(f'[make_line] ElevenLabs voice {args.elevenlabs} / {args.model}: {os.path.getsize(mp3)} bytes', file=sys.stderr)
    return mp3

def make_audio(args, base):
    aiff = base + '.aiff'
    if args.elevenlabs:
        src = elevenlabs_tts(args, base)
    elif args.audio:
        src = args.audio
    else:
        print('[make_line] WARNING: macOS `say` is a PLACEHOLDER voice — use --elevenlabs or --audio for anything shipped', file=sys.stderr)
        say = ['say', '-o', aiff, args.text]
        if args.voice: say[1:1] = ['-v', args.voice]
        run(say)
        src = aiff
    run(['ffmpeg', '-y', '-loglevel', 'error', '-i', src, '-ac', '1', '-ar', '16000', '-sample_fmt', 's16', base + '.wav'])
    run(['ffmpeg', '-y', '-loglevel', 'error', '-i', src, '-vn', '-c:a', 'aac', '-b:a', '96k', base + '.m4a'])
    if os.path.exists(aiff): os.remove(aiff)

def align(wav_path, words):
    from pocketsphinx import Decoder
    with wave.open(wav_path, 'rb') as w:
        assert w.getframerate() == 16000 and w.getnchannels() == 1 and w.getsampwidth() == 2, 'need 16 kHz mono 16-bit'
        pcm = w.readframes(w.getnframes())
        duration = w.getnframes() / 16000
    dec = Decoder(samprate=16000, loglevel='ERROR')
    try:
        frate = int(dec.config['frate'])
    except Exception:
        frate = 100
    for word, phones in EXTRA_WORDS.items():
        dec.add_word(word, phones, True)
    oov = [w for w in words if dec.lookup_word(w) is None]
    if oov:
        raise SystemExit(f'words missing from the dictionary: {oov} - add pronunciations to EXTRA_WORDS in {__file__}')
    dec.set_align_text(' '.join(words))
    dec.start_utt(); dec.process_raw(pcm, False, True); dec.end_utt()
    dec.set_alignment()            # second pass: phone-level timing
    dec.start_utt(); dec.process_raw(pcm, False, True); dec.end_utt()
    phones = []
    for word in dec.get_alignment():
        for ph in word:
            phones.append((ph.name, ph.start / frate, (ph.start + ph.duration) / frate))
    return phones, duration

def to_cues(phones, duration):
    cues = []
    for name, start, end in phones:
        base = re.sub(r'\d', '', name).upper()
        shape = PHONE_SHAPE.get(base, PHONE_SHAPE.get(name, DEFAULT_SHAPE))
        if cues and cues[-1]['value'] == shape:
            cues[-1]['end'] = end
        else:
            cues.append({'start': start, 'end': end, 'value': shape})
    merged = []
    for c in cues:  # merge blips into their predecessor so the mouth doesn't flicker
        if merged and (c['end'] - c['start']) < MIN_CUE:
            merged[-1]['end'] = c['end']
        else:
            merged.append(c)
    if merged:
        merged[-1]['end'] = max(merged[-1]['end'], duration)
        if merged[-1]['value'] != 'X':
            merged.append({'start': merged[-1]['end'], 'end': duration, 'value': 'X'})
    for c in merged:
        c['start'] = round(c['start'], 3); c['end'] = round(c['end'], 3)
    return merged

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('id'); ap.add_argument('text')
    ap.add_argument('--voice', default=None, help='macOS voice for `say` (default: system voice)')
    ap.add_argument('--audio', default=None, help='use this audio file instead of `say` (any ffmpeg format)')
    ap.add_argument('--elevenlabs', default=None, metavar='VOICE_ID', help='ElevenLabs text-to-speech with this voice id (ELEVENLABS_API_KEY in the env)')
    ap.add_argument('--model', default='eleven_multilingual_v2', help='ElevenLabs model id (with --elevenlabs)')
    args = ap.parse_args()
    os.makedirs(OUT, exist_ok=True)
    base = os.path.join(OUT, args.id)
    with open(base + '.txt', 'w') as f: f.write(args.text)
    make_audio(args, base)
    words = normalize(args.text)
    cues, duration = None, None
    try:
        phones, duration = align(base + '.wav', words)
        cues = to_cues(phones, duration)
        with open(base + '.json', 'w') as f:
            json.dump({'metadata': {'soundFile': f'{args.id}.m4a', 'duration': round(duration, 3)}, 'mouthCues': cues}, f, indent=1)
    except SystemExit:
        raise
    except Exception as e:  # keep the audio; the game falls back to amplitude lip-flap
        print(f'[make_line] alignment failed ({e}); line registered WITHOUT cues', file=sys.stderr)
    if duration is None:
        duration = float(subprocess.check_output(['ffprobe', '-v', 'error', '-show_entries', 'format=duration', '-of', 'csv=p=0', base + '.m4a']))
    os.remove(base + '.wav')
    lines_path = os.path.join(OUT, 'lines.json')
    lines = json.load(open(lines_path)) if os.path.exists(lines_path) else []
    lines = [l for l in lines if l.get('id') != args.id]
    entry = {'id': args.id, 'text': args.text, 'audioUrl': f'/assets/speech/{args.id}.m4a', 'duration': round(duration, 2)}
    if cues: entry['cuesUrl'] = f'/assets/speech/{args.id}.json'
    lines.append(entry)
    json.dump(lines, open(lines_path, 'w'), indent=2)
    if cues:
        counts = {}
        for c in cues: counts[c['value']] = counts.get(c['value'], 0) + 1
        print(f"{args.id}: {duration:.2f} s, {len(cues)} cues {counts}; first: " + ' '.join(f"{c['value']}@{c['start']}" for c in cues[:14]))
    else:
        print(f'{args.id}: {duration:.2f} s, no cues')

if __name__ == '__main__':
    main()