import os
import sys
import torch
import whisper

if hasattr(sys.stdout, 'reconfigure'):
    sys.stdout.reconfigure(encoding='utf-8')

venv_scripts = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'venv', 'Scripts')
if venv_scripts not in os.environ.get('PATH', ''):
    os.environ['PATH'] = venv_scripts + os.pathsep + os.environ.get('PATH', '')

device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = whisper.load_model('base', device=device)

prompt = "Trận thủy chiến Rạch Gầm - Xoài Mút năm 1785. Nguyễn Huệ, Nguyễn Nhạc, Nguyễn Ánh, Chiêu Tăng, Chiêu Sương, tướng Trương Văn Đa, quân Tây Sơn, quân Xiêm, Gia Định, Quy Nhơn, Mỹ Tho, Chân Lạp, hỏa pháo, thuyền chiến, của ta."

video_path = sys.argv[1] if len(sys.argv) > 1 else r"content\videos\ep1.mp4"
res = model.transcribe(
    video_path,
    language='vi',
    initial_prompt=prompt,
    condition_on_previous_text=False,  # Prevents repetition loops!
    no_speech_threshold=0.6,
    compression_ratio_threshold=2.4
)

print(f"=== Transcription for {video_path} ===")
for seg in res['segments']:
    print(f"{seg['start']:.2f}s -> {seg['end']:.2f}s: {seg['text']}")