"""
napisy.py - napisy SRT po polsku z dowolnego nagrania lub filmu, lokalnie (faster-whisper).
Autor: Piotr Olszewski, promptowy.com. Licencja: MIT. Sprawdzone 3.10 i 17.10.2026: faster-whisper 1.2.1, Python 3.13, RTX 4090, Windows 11.

Instalacja (raz):
  pip install faster-whisper "av<19"                     (PyAV 19 z 29.09.2026 psuje faster-whisper 1.2.1)
  pip install nvidia-cublas-cu12 nvidia-cudnn-cu12       (tylko karta Nvidia)
  winget install Gyan.FFmpeg                             (tylko do --wypal)

Użycie:
  python napisy.py film.mp4                      -> film.srt obok pliku
  python napisy.py film.mp4 --model large-v3-turbo
  python napisy.py film.mp4 --wypal              -> dodatkowo film_napisy.mp4 z napisami na obrazie (ffmpeg)
  python napisy.py film.mp4 --cpu                -> bez karty graficznej (wolniej, model int8)

Zasady podziału (jak w telewizji i na Netfliksie): maks. 42 znaki w wierszu, maks. 2 wiersze,
napis od 1 do 6 sekund, cięcie najchętniej po kropce, przecinku albo przed spójnikiem.
"""
import argparse, os, sys, subprocess, time, json

ZN_W_WIERSZU, MAX_WIERSZY, MAX_CZAS, MIN_CZAS = 42, 2, 6.0, 1.0
SPOJNIKI = {'i', 'a', 'ale', 'że', 'bo', 'lecz', 'oraz', 'albo', 'czy', 'gdy', 'kiedy', 'który', 'która', 'które', 'więc', 'jednak'}

def cuda_dll():
    # Na Windowsie faster-whisper potrzebuje cuBLAS i cuDNN. Szukamy ich w pakietach pip nvidia-*-cu12
    # (pip install nvidia-cublas-cu12 nvidia-cudnn-cu12) albo w bibliotekach PyTorcha, jeśli jest zainstalowany.
    import glob, site
    katalogi = []
    for sp in site.getsitepackages() + [site.getusersitepackages()]:
        katalogi += glob.glob(os.path.join(sp, 'nvidia', '*', 'bin'))
        katalogi += glob.glob(os.path.join(sp, 'torch', 'lib'))
    for lib in katalogi:
        if os.path.isdir(lib):
            try:
                os.add_dll_directory(lib)
            except (OSError, AttributeError):
                pass
            os.environ['PATH'] = lib + os.pathsep + os.environ['PATH']

def t_srt(t):
    ms = int(round(t * 1000)); g, ms = divmod(ms, 3600000); m, ms = divmod(ms, 60000); s, ms = divmod(ms, 1000)
    return f'{g:02d}:{m:02d}:{s:02d},{ms:03d}'

def zawin(slowa):
    """Dzieli listę słów na maks. 2 wiersze, najlepiej w połowie i nie przed krótkim słowem."""
    tekst = ' '.join(slowa)
    if len(tekst) <= ZN_W_WIERSZU: return tekst
    best, bi = 1e9, 1
    for i in range(1, len(slowa)):
        a, b = ' '.join(slowa[:i]), ' '.join(slowa[i:])
        # wiersz dłuższy niż 42 znaki to duża kara, ale nie wykluczenie (inaczej przy długim słowie podział się psuje)
        kara = 1000 * max(0, max(len(a), len(b)) - ZN_W_WIERSZU) + abs(len(a) - len(b)) + (15 if len(slowa[i-1]) <= 2 else 0) - (10 if slowa[i-1][-1] in ',.;:?!' else 0)
        if kara < best: best, bi = kara, i
    return ' '.join(slowa[:bi]) + '\n' + ' '.join(slowa[bi:])

def bloki(words):
    """Grupuje słowa (z czasami) w napisy."""
    out, cur = [], []
    def zamknij():
        if cur: out.append([cur[0][0], cur[-1][1], [w for _, _, w in cur]]); cur.clear()
    for (a, b, w) in words:
        if cur:
            tekst = ' '.join(x[2] for x in cur + [(a, b, w)])
            za_dlugo = len(tekst) > ZN_W_WIERSZU * MAX_WIERSZY - 6 or (b - cur[0][0]) > MAX_CZAS
            pauza = a - cur[-1][1] > 0.8
            if za_dlugo or pauza:
                # cofnij się do ostatniego dobrego miejsca cięcia (interpunkcja albo przed spójnikiem)
                cut = None
                for j in range(len(cur) - 1, 0, -1):
                    if cur[j-1][2][-1] in ',.;:?!' or cur[j][2].lower().strip(',.') in SPOJNIKI:
                        cut = j; break
                if za_dlugo and cut and cut >= len(cur) // 2:
                    reszta = cur[cut:]; del cur[cut:]; zamknij(); cur.extend(reszta)
                else:
                    zamknij()
        cur.append((a, b, w))
        if w[-1] in '.?!' and (cur[-1][1] - cur[0][0]) >= 1.5:
            zamknij()
    zamknij()
    # minimalny czas wyświetlania, bez nachodzenia na następny napis
    for i, n in enumerate(out):
        if n[1] - n[0] < MIN_CZAS:
            nast = out[i+1][0] - 0.05 if i + 1 < len(out) else n[0] + MIN_CZAS
            n[1] = max(n[1], min(n[0] + MIN_CZAS, nast))
    return out

def main():
    ap = argparse.ArgumentParser(description='Napisy SRT po polsku lokalnie (faster-whisper)')
    ap.add_argument('plik'); ap.add_argument('--model', default='large-v3')
    ap.add_argument('--jezyk', default='pl'); ap.add_argument('--cpu', action='store_true')
    ap.add_argument('--wypal', action='store_true', help='wypal napisy w obraz (wymaga ffmpeg)')
    ap.add_argument('--slowa', default='', help='podpowiedź: nazwiska i trudne słowa, np. "Falconbridge, Skawiński"')
    x = ap.parse_args()
    if not x.cpu: cuda_dll()
    from faster_whisper import WhisperModel
    t0 = time.time()
    m = WhisperModel(x.model, device='cpu' if x.cpu else 'cuda', compute_type='int8' if x.cpu else 'float16')
    t1 = time.time()
    segs, info = m.transcribe(x.plik, language=x.jezyk, beam_size=5, vad_filter=True, word_timestamps=True,
                              condition_on_previous_text=False, initial_prompt=x.slowa or None)
    words = []
    for s in segs:
        for i, w in enumerate(s.words):
            t = w.word.strip()
            if not t: continue
            if words and i > 0 and not w.word.startswith(' '):
                # kawałek bez spacji (np. „.pl”, „-go”) doklejamy do poprzedniego słowa
                a, _, p = words[-1]; words[-1] = (a, w.end, p + t)
            else:
                words.append((w.start, w.end, t))
    t2 = time.time()
    nap = bloki(words)
    srt = os.path.splitext(x.plik)[0] + '.srt'
    with open(srt, 'w', encoding='utf-8') as f:
        for i, (a, b, sl) in enumerate(nap, 1):
            f.write(f'{i}\n{t_srt(a)} --> {t_srt(b)}\n{zawin(sl)}\n\n')
    print(f'Gotowe: {srt} | {len(nap)} napisów | nagranie {info.duration/60:.1f} min | model {x.model} '
          f'wczytany w {t1-t0:.0f} s | rozpoznawanie {t2-t1:.0f} s')
    if x.wypal:
        wyj = os.path.splitext(x.plik)[0] + '_napisy.mp4'
        styl = 'FontName=Arial,FontSize=22,Outline=2,Shadow=0,MarginV=40'
        sciezka = srt.replace(os.sep, '/').replace(':', r'\:')  # ścieżka w formacie filtra ffmpeg
        subprocess.run(['ffmpeg', '-y', '-loglevel', 'error', '-i', x.plik, '-vf', f"subtitles='{sciezka}':force_style='{styl}'",
                        '-c:a', 'copy', wyj], check=True)
        print('Film z napisami:', wyj)

if __name__ == '__main__':
    main()