#!/usr/bin/env python3
"""Faster-Whisper: фактический текст и временная метка каждого слова."""
from __future__ import annotations
import argparse, csv, ctypes, json, os, re, shutil, subprocess, tempfile
from pathlib import Path

def srt_ts(t:float)->str:
    ms=round(t*1000); h,ms=divmod(ms,3600000); m,ms=divmod(ms,60000); s,ms=divmod(ms,1000)
    return f'{h:02d}:{m:02d}:{s:02d},{ms:03d}'

def cuda_runtime_ready()->tuple[bool,str]:
    if not shutil.which('nvidia-smi'):
        return False,'nvidia-smi not found'
    if os.name=='nt':
        try:
            ctypes.WinDLL('cublas64_12.dll')
        except OSError:
            return False,'cublas64_12.dll not found'
    return True,'ok'

def main():
    ap=argparse.ArgumentParser(); ap.add_argument('media'); ap.add_argument('--out',required=True)
    ap.add_argument('--model',default='small'); ap.add_argument('--language',default='ru')
    ap.add_argument('--device',default='auto'); ap.add_argument('--compute-type',default='auto')
    ap.add_argument('--glossary'); ap.add_argument('--beam-size',type=int,default=5)
    args=ap.parse_args()
    try: from faster_whisper import WhisperModel
    except ImportError: raise SystemExit('Установи faster-whisper из requirements-optional.txt')
    media=Path(args.media); prefix=Path(args.out); prefix.parent.mkdir(parents=True,exist_ok=True)
    glossary=''
    if args.glossary and Path(args.glossary).exists():
        glossary=Path(args.glossary).read_text(encoding='utf-8').strip()
    device=args.device; compute=args.compute_type
    cuda_ok,cuda_reason=cuda_runtime_ready()
    if device=='auto':
        device='cuda' if cuda_ok else 'cpu'
    elif device=='cuda' and not cuda_ok:
        raise SystemExit(
            f'CUDA preflight failed: {cuda_reason}. '
            'Install the matching CUDA 12 runtime or use --device cpu --compute-type int8.'
        )
    if compute=='auto': compute='int8' if device in ('cuda','cpu') else 'default'
    with tempfile.TemporaryDirectory() as td:
        wav=Path(td)/'speech.wav'
        q=subprocess.run(['ffmpeg','-y','-v','error','-i',str(media),'-vn','-ac','1','-ar','16000','-c:a','pcm_s16le',str(wav)])
        if q.returncode: raise SystemExit('FFmpeg не смог извлечь звук.')
        model=WhisperModel(args.model,device=device,compute_type=compute)
        seg_iter,info=model.transcribe(str(wav),language=args.language,beam_size=args.beam_size,temperature=0.0,
            word_timestamps=True,vad_filter=True,condition_on_previous_text=False,
            initial_prompt=(glossary or None),vad_parameters={'min_silence_duration_ms':180,'speech_pad_ms':100})
        segments=[]; words=[]
        for si,s in enumerate(seg_iter):
            sw=[]
            for wi,w in enumerate(s.words or []):
                token=(w.word or '').strip()
                if not token: continue
                row={'id':len(words),'segment_id':si,'word':token,'start':round(float(w.start),3),'end':round(float(w.end),3),
                     'probability':round(float(getattr(w,'probability',0.0) or 0.0),4)}
                words.append(row); sw.append(row['id'])
            segments.append({'id':si,'start':round(float(s.start),3),'end':round(float(s.end),3),'text':s.text.strip(),
                             'avg_logprob':float(getattr(s,'avg_logprob',0.0) or 0.0),'no_speech_prob':float(getattr(s,'no_speech_prob',0.0) or 0.0),'word_ids':sw})
    payload={'source':str(media),'model':args.model,'device':device,'compute_type':compute,
             'language':info.language,'language_probability':info.language_probability,
             'duration':float(getattr(info,'duration',0) or 0),'words':words,'segments':segments}
    (prefix.with_suffix('.words.json')).write_text(json.dumps(payload,ensure_ascii=False,indent=2),encoding='utf-8')
    (prefix.with_suffix('.segments.json')).write_text(json.dumps(segments,ensure_ascii=False,indent=2),encoding='utf-8')
    transcript=' '.join(w['word'] for w in words).replace(' ,',',').replace(' .','.').replace(' !','!').replace(' ?','?')
    (prefix.with_suffix('.actual_transcript.txt')).write_text(transcript+'\n',encoding='utf-8')
    with prefix.with_suffix('.words.csv').open('w',newline='',encoding='utf-8-sig') as f:
        wr=csv.DictWriter(f,fieldnames=['id','segment_id','word','start','end','probability']); wr.writeheader(); wr.writerows(words)
    srt=[]
    for i,s in enumerate(segments,1): srt += [str(i),f"{srt_ts(s['start'])} --> {srt_ts(s['end'])}",s['text'],'']
    prefix.with_suffix('.raw.srt').write_text('\n'.join(srt),encoding='utf-8')
    print(json.dumps({'words':len(words),'segments':len(segments),'device':device,'compute_type':compute,
                      'first':words[:5],'last':words[-5:]},ensure_ascii=False,indent=2))

if __name__=='__main__': main()