import os
import subprocess
import tempfile
from pathlib import Path
from typing import Optional

import streamlit as st
from faster_whisper import WhisperModel
from googletrans import Translator

# ---------------- CONFIG ----------------
MODEL_NAME = os.environ.get("WHISPER_MODEL", "large-v2")
DEVICE = "cuda" if (os.environ.get("CUDA_VISIBLE_DEVICES") or False) else "cpu"

_model = None
def get_model():
    global _model
    if _model is None:
        compute_type = "float16" if DEVICE.startswith("cuda") else "int8"
        _model = WhisperModel(MODEL_NAME, device=DEVICE, compute_type=compute_type)
    return _model


def extract_audio(input_video_path: str, output_audio_path: str):
    cmd = [
        "ffmpeg", "-y", "-i", input_video_path,
        "-vn", "-acodec", "pcm_s16le", "-ar", "16000", "-ac", "1",
        output_audio_path,
    ]
    subprocess.run(cmd, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)


def segments_to_srt(segments):
    def fmt_time(s):
        h = int(s // 3600)
        m = int((s % 3600) // 60)
        sec = s % 60
        return f"{h:02d}:{m:02d}:{sec:06.3f}".replace('.', ',')

    srt_lines = []
    for i, seg in enumerate(segments, start=1):
        start = fmt_time(seg["start"])
        end = fmt_time(seg["end"])
        text = seg["text"].strip()
        srt_lines.append(f"{i}\n{start} --> {end}\n{text}\n")
    return "\n".join(srt_lines)


def transcribe_and_translate(video_file: str, target_lang: Optional[str], burn_subs: bool):
    model = get_model()
    tempdir = Path(tempfile.mkdtemp())
    input_path = Path(video_file)
    audio_path = tempdir / "audio.wav"
    srt_path = tempdir / f"subtitles_{input_path.stem}.srt"
    processed_video_path = None

    extract_audio(str(input_path), str(audio_path))

    segments_all = []
    transcribe_options = {"beam_size": 5, "word_timestamps": False}
    for segment in model.transcribe(str(audio_path), beam_size=5, vad_filter=True, **transcribe_options):
        segments_all.append({"start": segment.start, "end": segment.end, "text": segment.text})

    if target_lang and target_lang.lower() not in ["", "none"]:
        translator = Translator()
        translated_segments = []
        for seg in segments_all:
            src_text = seg["text"].strip()
            try:
                res = translator.translate(src_text, dest=target_lang)
                translated_text = res.text
            except Exception:
                translated_text = src_text
            translated_segments.append({"start": seg["start"], "end": seg["end"], "text": translated_text})
        segments_used = translated_segments
    else:
        segments_used = segments_all

    srt_text = segments_to_srt(segments_used)
    srt_path.write_text(srt_text, encoding="utf-8")

    if burn_subs:
        out_video = tempdir / f"burned_{input_path.name}"
        cmd = [
            "ffmpeg", "-y", "-i", str(input_path),
            "-vf", f"subtitles={str(srt_path)}:force_style='FontName=Arial,FontSize=24'",
            "-c:a", "copy", str(out_video),
        ]
        subprocess.run(cmd, check=True)
        processed_video_path = str(out_video)

    return str(srt_path), processed_video_path


# ---------------- UI (Streamlit) ----------------
st.set_page_config(page_title="Video Subtitle Editor + Translator", layout="wide")
st.title("🎬 Video Subtitle Editor + Translator (Streamlit)")

video_file = st.file_uploader("Upload your video (mp4, mov, mkv)", type=["mp4", "mov", "mkv"])
lang_choice = st.selectbox(
    "Translate subtitles to:",
    ["None", "English (en)", "Urdu (ur)", "Hindi (hi)", "Spanish (es)", "French (fr)", "German (de)"]
)
burn_option = st.checkbox("Burn subtitles into video (hardcoded) - slow but permanent", value=False)

if st.button("Run"):
    if video_file is None:
        st.warning("Please upload a video file first.")
    else:
        with st.spinner("Processing... Please wait ⏳"):
            temp_input = Path(tempfile.mkdtemp()) / video_file.name
            with open(temp_input, "wb") as f:
                f.write(video_file.read())

            lang_code = lang_choice.split("(")[-1].replace(")", "").strip().lower() if "(" in lang_choice else "none"
            try:
                srt_path, processed_video = transcribe_and_translate(str(temp_input), lang_code, burn_option)
                st.success("✅ Done! Subtitles generated successfully.")
                st.download_button("📥 Download SRT", open(srt_path, "rb"), file_name=os.path.basename(srt_path))
                if processed_video:
                    st.download_button("📽️ Download Video with Subtitles", open(processed_video, "rb"), file_name=os.path.basename(processed_video))
            except subprocess.CalledProcessError as e:
                st.error(f"ffmpeg error: {e}")
            except Exception as e:
                st.error(f"Error: {e}")