#!/usr/bin/env python3
"""narration.md -> digest-audio.mp3 — versione resilient:
- cache su disco dei segmenti WAV (riparte da dove era: i segmenti gia' fatti non si rifanno)
- retry con backoff esponenziale sui 429/503 (fino a 8 tentativi, attesa max ~5 min)
- pausa tra le chiamate per stare lontano dal rate limit
- guardia anti-divagazione (durata attesa vs reale)
Uso: python3 narrate_to_audio.py [narration.md] [output.mp3]"""
import base64, hashlib, json, os, re, subprocess, sys, time, urllib.request, urllib.error, wave

BASEDIR = '/opt/data/hn-oss-digest'
C = json.load(open(BASEDIR + '/creds.json'))
SRC = sys.argv[1] if len(sys.argv) > 1 else BASEDIR + '/narration.md'
OUT = sys.argv[2] if len(sys.argv) > 2 else BASEDIR + '/digest-audio.mp3'
CACHE = os.path.join(BASEDIR, 'tts-cache')
os.makedirs(CACHE, exist_ok=True)
STYLE = ('Speak in Italian with a native Italian accent, warm and passionate like a tech '
         'podcast narrator, conversational pace, natural pauses between sentences, a hint '
         'of wonder on curious details, friendly closing energy: \n')
PAUSE_BETWEEN = 12   # secondi tra una chiamata TTS e l'altra

def segments(text, maxlen=1500):
    paras = [p.strip() for p in re.split(r'\n\s*\n', text) if p.strip()]
    segs, cur = [], ''
    for p in paras:
        if len(cur) + len(p) + 2 <= maxlen:
            cur = (cur + '\n\n' + p).strip()
        else:
            if cur:
                segs.append(cur)
            while len(p) > maxlen:
                cut = p.rfind('. ', 0, maxlen)
                cut = cut + 1 if cut > 200 else maxlen
                segs.append(p[:cut].strip())
                p = p[cut:].strip()
            cur = p
    if cur:
        segs.append(cur)
    return segs

MODELS = ['gemini-3.8-flash-lite-tts', 'gemini-3.8-flash-tts',
          'gemini-3.1-flash-tts-preview', 'gemini-2.5-flash-preview-tts']

def tts_once(text, model):
    payload = {
        'contents': [{'parts': [{'text': STYLE + text}]}],
        'generationConfig': {
            'responseModalities': ['AUDIO'],
            'speechConfig': {'voiceConfig': {'prebuiltVoiceConfig': {'voiceName': 'Kore'}}},
        },
    }
    url = ('https://generativelanguage.googleapis.com/v1beta/models/' + model + ':generateContent?key='
           + C['gemini_key'])
    req = urllib.request.Request(url, data=json.dumps(payload).encode(),
                                 headers={'Content-Type': 'application/json'})
    with urllib.request.urlopen(req, timeout=300) as r:
        x = json.loads(r.read().decode())
    cand = x['candidates'][0]
    part = cand['content']['parts'][0]
    data = base64.b64decode(part['inlineData']['data'])
    rate = 24000
    for tok in part['inlineData'].get('mimeType', '').split(';'):
        if tok.strip().startswith('rate='):
            rate = int(tok.split('=')[1])
    if cand.get('finishReason') != 'STOP':
        raise RuntimeError('finishReason=' + str(cand.get('finishReason')))
    return data, rate

def tts(text):
    """Cache + cascata di modelli + backoff: se un modello e' a secco (429) prova il prossimo."""
    key = hashlib.sha256((STYLE + text).encode()).hexdigest()[:16]
    cache_file = os.path.join(CACHE, key + '.json')
    if os.path.exists(cache_file):
        d = json.load(open(cache_file))
        print('  [cache]', end=' ')
        return base64.b64decode(d['data']), d['rate']
    models = [m for m in MODELS]
    delay = 60
    for cycle in range(3):
        for model in list(models):
            try:
                data, rate = tts_once(text, model)
                json.dump({'data': base64.b64encode(data).decode(), 'rate': rate}, open(cache_file, 'w'))
                return data, rate
            except urllib.error.HTTPError as e:
                if e.code == 429:
                    print(f'  [{model.split("-tts")[0]}: esaurito]', end=' ', flush=True)
                    continue
                if e.code in (500, 503):
                    print(f'  [{model}: {e.code}, riprovo]', end=' ', flush=True)
                    time.sleep(30)
                    continue
                raise
        if cycle < 2:
            print(f'  [tutti a secco, attendo {delay}s]', end=' ', flush=True)
            time.sleep(delay)
            delay *= 3
    raise RuntimeError('tts: tutti i modelli esauriti')

text = open(SRC, encoding='utf-8').read()
segs = segments(text)
print(f'narrazione: {len(text)} char -> {len(segs)} segmenti audio')

frames, rate = [], 24000
for i, s in enumerate(segs, 1):
    data, r = tts(s)
    dur = len(data) / (r * 2)
    atteso = max(6.0, len(s) / 14.0)
    if dur > atteso * 1.8:
        print(f'  segmento {i}: {dur:.0f}s vs attesi ~{atteso:.0f}s -> retry (senza cache)')
        os.remove(os.path.join(CACHE, hashlib.sha256((STYLE + s).encode()).hexdigest()[:16] + '.json'))
        data, r = tts(s)
        dur = len(data) / (r * 2)
        if dur > atteso * 1.8:
            print(f'  segmento {i}: ancora {dur:.0f}s, lo tengo (possibile divagazione)')
    rate = r
    frames.append(data)
    print(f'  segmento {i}/{len(segs)}: {dur:.1f}s')
    if i < len(segs):
        time.sleep(PAUSE_BETWEEN)

wav = OUT.replace('.mp3', '.wav')
with wave.open(wav, 'wb') as w:
    w.setnchannels(1)
    w.setsampwidth(2)
    w.setframerate(rate)
    for f in frames:
        w.writeframes(f)

subprocess.run(['ffmpeg', '-y', '-loglevel', 'error', '-i', wav,
                '-codec:a', 'libmp3lame', '-b:a', '64k', '-ar', '24000', '-ac', '1', OUT],
               check=True)
print('OK:', OUT, os.path.getsize(OUT) // 1024, 'KB')