#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""用 faster-whisper 对翻唱音频做中文转写，输出句子级 + 词级时间戳。

用法: python3 transcribe.py [音频wav] [模型名,默认large-v3]
输出: 00-音源/九九八十一-whisper.json  00-音源/九九八十一-whisper.srt  00-音源/九九八十一-whisper.txt
"""
import json, os, sys, time

ROOT = os.path.dirname(os.path.abspath(__file__))
os.chdir(ROOT)
AUDIO = sys.argv[1] if len(sys.argv) > 1 else "00-音源/九九八十一-咻咻满翻唱.wav"
MODEL = sys.argv[2] if len(sys.argv) > 2 else "large-v3"

from faster_whisper import WhisperModel

lyr = open("00-音源/lyrics-九九八十一.txt", encoding="utf-8").read().replace("\n", " ")
prompt = lyr[:400]

print(f"[i] 加载模型 {MODEL} (cpu/int8)…", flush=True)
t0 = time.time()
model = WhisperModel(MODEL, device="cpu", compute_type="int8")
print(f"[i] 模型就绪 {time.time()-t0:.0f}s，开始转写 {AUDIO}", flush=True)

t0 = time.time()
segments, info = model.transcribe(
    AUDIO, language="zh",
    beam_size=int(os.environ.get("WHISPER_BEAM", "5")),
    initial_prompt=prompt,
    word_timestamps=True,
    vad_filter=os.environ.get("WHISPER_VAD", "0") == "1",
    condition_on_previous_text=False,
)
segs = []
for s in segments:
    words = [{"w": w.word, "s": round(w.start, 2), "e": round(w.end, 2)} for w in (s.words or [])]
    segs.append({"start": round(s.start, 2), "end": round(s.end, 2), "text": s.text.strip(), "words": words})
    print(f"{s.start:7.2f} - {s.end:7.2f}  {s.text.strip()}", flush=True)
print(f"[i] 转写完成 {time.time()-t0:.0f}s，共 {len(segs)} 段", flush=True)

with open("00-音源/九九八十一-whisper.json", "w", encoding="utf-8") as f:
    json.dump({"audio": AUDIO, "model": MODEL, "duration": info.duration, "segments": segs}, f, ensure_ascii=False, indent=1)


def ts(t):
    h = int(t // 3600); m = int(t % 3600 // 60); s = t % 60
    return f"{h:02d}:{m:02d}:{s:06.3f}".replace(".", ",")


with open("00-音源/九九八十一-whisper.srt", "w", encoding="utf-8") as f:
    for i, s in enumerate(segs, 1):
        f.write(f"{i}\n{ts(s['start'])} --> {ts(s['end'])}\n{s['text']}\n\n")

with open("00-音源/九九八十一-whisper.txt", "w", encoding="utf-8") as f:
    for s in segs:
        f.write(f"{s['start']:7.2f}\t{s['end']:7.2f}\t{s['text']}\n")
print("OK")