#!/usr/bin/env python3
from pathlib import Path
import json, subprocess, time
from faster_whisper import WhisperModel

ROOT = Path('/home/spazio/new_audio_transcriptions_2026-10-20')
SRC = ROOT / 'audio'
OUT = ROOT / 'out'
OUT.mkdir(parents=True, exist_ok=True)
model = WhisperModel('small', device='cpu', compute_type='int8')
media = sorted(SRC.glob('*.mp3'))
for n, p in enumerate(media, 1):
    txt = OUT / f'{p.stem}.txt'
    srt = OUT / f'{p.stem}.srt'
    meta = OUT / f'{p.stem}.json'
    if txt.exists() and srt.exists() and meta.exists():
        try:
            m = json.loads(meta.read_text('utf8'))
            if m.get('complete'):
                print(f'SKIP {n}/{len(media)} {p.name}', flush=True)
                continue
        except Exception:
            pass
    print(f'START {n}/{len(media)} {p.name}', flush=True)
    started = time.time()
    seg_iter, info = model.transcribe(
        str(p), language='ru', vad_filter=True, beam_size=5,
        condition_on_previous_text=True, word_timestamps=False
    )
    segments = []
    for i, seg in enumerate(seg_iter, 1):
        text = seg.text.strip()
        if text:
            segments.append({'start': float(seg.start), 'end': float(seg.end), 'text': text})
        if i % 50 == 0:
            checkpoint = OUT / f'{p.stem}.checkpoint.txt'
            checkpoint.write_text('\n\n'.join(x['text'] for x in segments) + '\n', 'utf8')
            print(f'PROGRESS {n}/{len(media)} segments={i} end={segments[-1]["end"]:.1f}', flush=True)
    def ts(t):
        ms = int(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}'
    txt.write_text('\n\n'.join(x['text'] for x in segments).strip() + '\n', 'utf8')
    srt.write_text('\n\n'.join(f'{i}\n{ts(x["start"])} --> {ts(x["end"])}\n{x["text"]}' for i, x in enumerate(segments, 1)) + '\n', 'utf8')
    probe = subprocess.run(['ffprobe','-v','error','-show_entries','format=duration','-of','json',str(p)], capture_output=True, text=True, check=True)
    duration = float(json.loads(probe.stdout)['format']['duration'])
    last_end = segments[-1]['end'] if segments else 0.0
    meta.write_text(json.dumps({
        'source': p.name, 'language': getattr(info, 'language', 'ru'),
        'duration': duration, 'last_speech_end': last_end,
        'trailing_gap': duration-last_end, 'segments': len(segments),
        'complete': True, 'elapsed': time.time()-started
    }, ensure_ascii=False, indent=2), 'utf8')
    (OUT / f'{p.stem}.checkpoint.txt').unlink(missing_ok=True)
    print(f'DONE {n}/{len(media)} {p.name} duration={duration:.1f} last={last_end:.1f} gap={duration-last_end:.1f}', flush=True)
print(f'BATCH_DONE files={len(media)}', flush=True)