import streamlit as st
import whisper
import tempfile
from pathlib import Path
import torch


def format_timestamp(seconds: float) -> str:
    """Converte segundos em timestamp HH:MM:SS.mmm para legenda."""
    ms = int((seconds - int(seconds)) * 1000)
    seconds = int(seconds)
    s = seconds % 60
    minutes = (seconds // 60) % 60
    hours = seconds // 3600
    return f"{hours:02d}:{minutes:02d}:{s:02d}.{ms:03d}"


def segments_to_vtt(segments) -> str:
    """Converte segments do Whisper em arquivo WEBVTT."""
    lines = ["WEBVTT", ""]
    for idx, seg in enumerate(segments, start=1):
        start = format_timestamp(seg["start"])
        end = format_timestamp(seg["end"])
        text = seg["text"].strip()
        lines.append(f"{idx}")
        lines.append(f"{start} --> {end}")
        lines.append(text)
        lines.append("")
    return "\n".join(lines)


# Configuração básica da página
st.set_page_config(
    page_title="FFKeiro – Áudio → Texto",
    page_icon="🎧",
    layout="centered",
)

st.title("🎧 FFKeiro – Conversor de Áudio para Texto")
st.write(
    "Envie um arquivo de áudio para gerar a transcrição em **texto** "
    "e uma legenda sincronizada em **WEBVTT (.vtt)**."
)

# Sidebar com opções
model_size = st.sidebar.selectbox(
    "Modelo Whisper",
    options=["tiny", "base", "small", "medium", "large"],
    index=2,  # "small" como padrão
    help="Modelos menores são mais rápidos; modelos maiores tendem a ter melhor qualidade.",
)

language = st.sidebar.text_input(
    "Idioma (código ISO)",
    value="pt",
    help='Ex.: "pt" para português, "en" para inglês, etc.',
)

device = "cuda" if torch.cuda.is_available() else "cpu"
st.sidebar.write(f"Device detectado: **{device}**")

# Upload de áudio
uploaded_file = st.file_uploader(
    "Envie o arquivo de áudio",
    type=["wav", "mp3", "m4a", "ogg", "flac"],
)

if uploaded_file is not None:
    st.audio(uploaded_file)
    st.write("Arquivo recebido:", uploaded_file.name)

    if st.button("🔁 Transcrever áudio"):
        with st.spinner("Carregando modelo Whisper e transcrevendo..."):
            # Carrega o modelo
            model = whisper.load_model(model_size, device=device)

            # Salva o arquivo enviado em um temp file
            with tempfile.NamedTemporaryFile(
                delete=False,
                suffix=Path(uploaded_file.name).suffix,
            ) as tmp:
                tmp.write(uploaded_file.read())
                tmp_path = tmp.name

            # Transcreve
            result = model.transcribe(tmp_path, language=language)
            text = result.get("text", "").strip()
            segments = result.get("segments", [])

        if not text:
            st.error("Não foi possível obter transcrição.")
        else:
            st.success("Transcrição concluída!")

            # Mostra texto
            st.subheader("📝 Transcrição (Texto)")
            st.text_area("Texto completo", value=text, height=250)

            # Gera arquivos para download
            txt_bytes = text.encode("utf-8")
            vtt_content = segments_to_vtt(segments) if segments else ""
            vtt_bytes = vtt_content.encode("utf-8") if vtt_content else b""

            col1, col2 = st.columns(2)
            with col1:
                st.download_button(
                    label="⬇️ Baixar TXT",
                    data=txt_bytes,
                    file_name=Path(uploaded_file.name).stem + ".txt",
                    mime="text/plain",
                )
            with col2:
                if vtt_bytes:
                    st.download_button(
                        label="⬇️ Baixar legenda VTT",
                        data=vtt_bytes,
                        file_name=Path(uploaded_file.name).stem + ".vtt",
                        mime="text/vtt",
                    )
                else:
                    st.info("Legenda VTT não disponível.")