import os
import uuid
import pandas as pd
import gradio as gr

from faster_whisper import WhisperModel
from huggingface_hub import HfApi, upload_file

# =========================
# ENV
# =========================
HF_TOKEN = os.getenv("HF_TOKEN")
DATASET_REPO = os.getenv("DATASET_REPO")

# =========================
# LOAD MODEL
# =========================
model = WhisperModel("base")

# =========================
# TRANSCRIBE FUNCTION
# =========================
def transcribe_video(video_file):

    if video_file is None:
        return "No file uploaded.", None

    file_id = str(uuid.uuid4())[:8]

    segments, info = model.transcribe(
        video_file,
        beam_size=5
    )

    transcript = ""

    for segment in segments:
        transcript += f"[{segment.start:.2f} - {segment.end:.2f}] {segment.text}\n"

    # =========================
    # SAVE LOCAL TXT
    # =========================
    txt_name = f"transcript_{file_id}.txt"

    with open(txt_name, "w", encoding="utf-8") as f:
        f.write(transcript)

    # =========================
    # SAVE CSV METADATA
    # =========================
    csv_name = f"metadata_{file_id}.csv"

    df = pd.DataFrame([
        {
            "file_id": file_id,
            "original_file": os.path.basename(video_file),
            "transcript": transcript
        }
    ])

    df.to_csv(csv_name, index=False)

    # =========================
    # UPLOAD TO DATASET
    # =========================
    api = HfApi(token=HF_TOKEN)

    upload_file(
        path_or_fileobj=txt_name,
        path_in_repo=f"transcripts/{txt_name}",
        repo_id=DATASET_REPO,
        repo_type="dataset",
        token=HF_TOKEN,
    )

    upload_file(
        path_or_fileobj=csv_name,
        path_in_repo=f"metadata/{csv_name}",
        repo_id=DATASET_REPO,
        repo_type="dataset",
        token=HF_TOKEN,
    )

    return transcript, txt_name

# =========================
# UI
# =========================
app = gr.Interface(
    fn=transcribe_video,
    inputs=gr.Video(label="Upload Video"),
    outputs=[
        gr.Textbox(label="Transcript"),
        gr.File(label="Download TXT")
    ],
    title="AI Video Transcriber",
    description="Chinese + English AI transcription"
)

app.launch()