from faster_whisper import WhisperModel
import sys
import os

os.environ["HF_HOME"] = "/var/www/html/whisper_cache"
os.environ["HUGGINGFACE_HUB_CACHE"] = "/var/www/html/whisper_cache"
os.environ["XDG_CACHE_HOME"] = "/var/www/html/whisper_cache"
os.environ["HF_HUB_DISABLE_XET"] = "1"

def format_vtt_time(seconds):
    minutes = int(seconds // 60)
    seconds = seconds % 60

    return f"{minutes:02}:{seconds:06.3f}"


if len(sys.argv) < 3:
    print("Uso: python transcribe.py <video> <output_dir>")
    sys.exit(1)

video_file = sys.argv[1]
output_dir = sys.argv[2]

if not os.path.exists(video_file):
    print(f"Arquivo não encontrado: {video_file}")
    sys.exit(1)

os.makedirs(output_dir, exist_ok=True)

base_name = os.path.splitext(
    os.path.basename(video_file)
)[0]

vtt_file = os.path.join(
    output_dir,
    f"{base_name}.vtt"
)

print("Carregando modelo...")

model = WhisperModel(
    "base",
    device="cpu",
    compute_type="int8",
    download_root="/var/www/html/whisper_cache"
)

print("Iniciando transcrição...")

segments, info = model.transcribe(
    video_file,
    language="pt",
    beam_size=5,
    vad_filter=True,
    word_timestamps=False
)

with open(
    vtt_file,
    "w",
    encoding="utf-8"
) as file:

    file.write("WEBVTT\n\n")

    for segment in segments:

        start = format_vtt_time(
            segment.start
        )

        end = format_vtt_time(
            segment.end
        )

        text = segment.text.strip()

        file.write(
            f"{start} --> {end}\n"
        )

        file.write(
            f"{text}\n\n"
        )

print(
    f"Legenda criada com sucesso: {vtt_file}"
)