#!/usr/bin/env python3
"""STT + Speaker Diarization pipeline (frigate.vpn).
ASR: whisper.cpp (Vulkan) | Diarization: pyannote Community-1 (CPU)
Uso: transcribe.py <audio> [--lang it] [--model large-v3-turbo] [--out base]
     [--prompt \"frase personale\"] [--glossary glossario.json]
"""
import argparse, json, os, re, subprocess, sys, tempfile

WHISPER_CLI = os.environ.get('WHISPER_CLI', os.path.expanduser('~/dev/whisper.cpp/build/bin/whisper-cli'))
MODELS_DIR  = os.environ.get('MODELS_DIR', os.path.expanduser('~/dev/whisper.cpp/models'))
PYTHON      = os.environ.get('STT_PYTHON', os.path.expanduser('~/dev/stt-venv/bin/python'))

# ---------------------------------------------------------------------------
# Glossario personale: estensione vocabolario + normalizzazione deterministica
# ---------------------------------------------------------------------------
# Formato (JSON):
# {
#   "people":         [{"canonical": "Giulia Bianchi", "aliases": ["Giulia Bianci", ...]}],
#   "organizations":  [{"canonical": "...", "aliases": [...]}],
#   "terms":          [{"canonical": "fine-tuning", "aliases": ["fine tuning", "fain tuning"]}]
# }
# La normalizzazione sostituisce gli alias con la forma canonica SOLO su match
# esatto (case-insensitive, confini di parola). Le sostituzioni sono annotate
# nel campo "raw" del segmento per audit/rollback.

def load_glossary(path_or_json):
    """Accetta un path JSON o una stringa JSON; restituisce la mappa
    alias->canonical (alias in lowercase). Invalido/vuoto -> {}."""
    if not path_or_json:
        return {}
    try:
        data = path_or_json
        if isinstance(data, str):
            stripped = data.strip()
            if stripped.startswith('{'):
                data = json.loads(stripped)          # stringa JSON
            else:
                with open(stripped, encoding='utf-8') as f:  # path
                    data = json.load(f)
        flat = {}
        for group in ('people', 'organizations', 'terms', 'custom'):
            for entry in (data.get(group) or []):
                canon = str(entry.get('canonical') or '').strip()
                if not canon:
                    continue
                for alias in (entry.get('aliases') or []):
                    alias = str(alias).strip()
                    if alias and alias.lower() != canon.lower():
                        flat[alias.lower()] = canon
        # Formato alternativo semplice: {"alias": "canonical"}
        if not flat:
            for alias, canon in data.items():
                if isinstance(canon, str) and isinstance(alias, str):
                    flat[alias.lower()] = canon
        return flat
    except Exception:
        return {}


def normalize_text(text, flat_map):
    """Sostituisce alias (case-insensitive, confini parola) con la forma
    canonica. Le sostituzioni avvengono su segnaposto univoci: un alias non
    può ri-matchare testo generato da un'altra sostituzione (no cascate).
    Restituisce (nuovo_testo, sostituzioni_fatte)."""
    if not flat_map or not text:
        return text, 0
    out, count = text, 0
    placeholders = []
    # Passo 1: alias -> segnaposto (alias più lunghi prima, evita match parziali)
    for alias in sorted(flat_map, key=len, reverse=True):
        canon = flat_map[alias]
        pattern = r'(?<!\w)' + re.escape(alias) + r'(?!\w)'
        ph = f'\x00PH{len(placeholders)}\x00'
        new_out, n = re.subn(pattern, ph, out, flags=re.IGNORECASE)
        if n:
            placeholders.append((ph, canon))
            count += n
            out = new_out
    # Passo 2: segnaposto -> forma canonica
    for ph, canon in placeholders:
        out = out.replace(ph, canon)
    return out, count


def asr(audio, model, lang, tmp, prompt=None):
    model_path = os.path.join(MODELS_DIR, f'ggml-{model}.bin')
    out_json = os.path.join(tmp, 'asr.json')
    vad_model = os.path.join(MODELS_DIR, 'ggml-silero-v6.2.0.bin')
    cmd = [WHISPER_CLI, '-m', model_path, '-f', audio, '-l', lang,
           '-oj', '-of', os.path.join(tmp, 'asr'), '--no-prints',
           '-vm', vad_model, '--vad',
           '--suppress-regex',
           r'(Grazie a tutti|Grazie per l.attenzione|Thank you|Thanks for watching|Sottotitoli creati|Sottotitoli)']
    # Initial prompt (bias vocabolario personale, nomi, gergo): boost
    # zero-training. --carry-initial-prompt lo riapplica su ogni finestra
    # (necessario per audio lunghi; verificare con whisper-cli -h).
    if prompt:
        cmd += ['--prompt', prompt, '--carry-initial-prompt']
    subprocess.run(cmd, check=True, capture_output=True)
    data = json.load(open(out_json))
    segs = []
    for s in data['transcription']:
        text = s['text'].strip()
        # Filtro finale: scarta segmenti che sono solo frasi allucinate
        # (rete di sicurezza oltre a VAD e suppress-regex)
        if text.lower() in {'grazie', 'grazie a tutti', 'thank you', 'thanks',
                            "grazie per l'attenzione", 'sottotitoli',
                            'sottotitoli creati', 'sottotitoli creati da'}:
            continue
        segs.append({'start': s['offsets']['from']/1000.0,
                     'end':   s['offsets']['to']/1000.0,
                     'text':  text})
    return segs


def apply_glossary(segs, flat_map):
    """Normalizza i testi dei segmenti con il glossario. I segmenti modificati
    ricevono 'raw' = testo ASR originale (audit/rollback)."""
    for seg in segs:
        new_text, count = normalize_text(seg['text'], flat_map)
        if count:
            seg['raw'] = seg['text']
            seg['text'] = new_text
    return segs


def diarize(audio, tmp):
    script = f'''
import json, sys
import numpy as np, soundfile as sf, torch
from pyannote.audio import Pipeline
data, sr = sf.read({audio!r}, dtype='float32')
wav = torch.from_numpy(data.T if data.ndim > 1 else data[None, :]).float()
p = Pipeline.from_pretrained('pyannote/speaker-diarization-community-1')
out = p({{'waveform': wav, 'sample_rate': sr}})
diar = out.speaker_diarization if hasattr(out, 'speaker_diarization') else out
turns = [{{'start': t.start, 'end': t.end, 'speaker': sp}}
         for t, _, sp in diar.itertracks(yield_label=True)]
json.dump(turns, open({os.path.join(tmp, 'diar.json')!r}, 'w'))
'''
    sp = os.path.join(tmp, 'diar.py')
    open(sp, 'w').write(script)
    subprocess.run([PYTHON, sp], check=True)
    return json.load(open(os.path.join(tmp, 'diar.json')))

def overlap(a, b):
    return max(0.0, min(a['end'], b['end']) - max(a['start'], b['start']))

def merge(segs, turns):
    for s in segs:
        best, best_ov = None, 0.0
        for t in turns:
            ov = overlap(s, t)
            if ov > best_ov:
                best, best_ov = t['speaker'], ov
        s['speaker'] = best if best_ov > 0.05 else 'UNKNOWN'
    return segs

def to_srt(segs, path):
    def ts(x):
        h, m = int(x//3600), int(x%3600//60)
        s, ms = int(x%60), int((x-int(x))*1000)
        return f'{h:02d}:{m:02d}:{s:02d},{ms:03d}'
    with open(path, 'w') as f:
        for i, s in enumerate(segs, 1):
            f.write(f'{i}\n{ts(s["start"])} --> {ts(s["end"])}\n')
            f.write(f'[{s["speaker"]}] {s["text"]}\n\n')

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('audio')
    ap.add_argument('--lang', default='it')
    ap.add_argument('--model', default='large-v3-turbo')
    ap.add_argument('--out', default=None)
    ap.add_argument('--prompt', default=None,
                    help='initial prompt personale (nomi, termini, gergo)')
    ap.add_argument('--glossary', default=None,
                    help='path a glossario JSON (alias -> forma canonica)')
    args = ap.parse_args()
    base = args.out or os.path.splitext(args.audio)[0]
    flat_map = load_glossary(args.glossary)
    with tempfile.TemporaryDirectory() as tmp:
        print(f'[1/3] ASR whisper.cpp ({args.model}, {args.lang})...', file=sys.stderr)
        segs = asr(args.audio, args.model, args.lang, tmp, prompt=args.prompt)
        print(f'      {len(segs)} segmenti', file=sys.stderr)
        if flat_map:
            segs = apply_glossary(segs, flat_map)
            print(f'      glossario: {len(flat_map)} alias applicati', file=sys.stderr)
        print('[2/3] Diarization pyannote Community-1...', file=sys.stderr)
        turns = diarize(args.audio, tmp)
        print(f'      {len(turns)} turni parlante', file=sys.stderr)
        print('[3/3] Merge ASR + diarization...', file=sys.stderr)
        segs = merge(segs, turns)
    out_json = base + '.json'
    out_srt  = base + '.srt'
    json.dump({'segments': segs}, open(out_json, 'w'), ensure_ascii=False, indent=2)
    to_srt(segs, out_srt)
    print(f'OK: {out_json} + {out_srt}', file=sys.stderr)
    for s in segs:
        print(f'{s["start"]:6.2f}-{s["end"]:6.2f} [{s["speaker"]}] {s["text"]}')

if __name__ == '__main__':
    main()