import os
import sys
import time
import torch
import whisper

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

# Ensure ffmpeg in PATH
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', '')

video_path = sys.argv[1] if len(sys.argv) > 1 else r"content\videos\ep1.mp4"
model_name = sys.argv[2] if len(sys.argv) > 2 else "base"

device = "cuda" if torch.cuda.is_available() else "cpu"
t0 = time.time()
print(f"Loading Whisper '{model_name}' on {device}...")
model = whisper.load_model(model_name, device=device)

print(f"Transcribing {video_path}...")
res = model.transcribe(video_path, language='vi')
t1 = time.time()

print(f"Completed in {t1 - t0:.2f}s!")
for seg in res['segments']:
    print(f"{seg['start']:.2f}s -> {seg['end']:.2f}s: {seg['text']}")