# -*- coding: utf-8 -*-

import whisper
import srt
from datetime import timedelta

def criar_legendas_curtas(audio_path, srt_path, max_chars=20):
    # 1. Carrega o modelo e transcreve com foco nas palavras
    print("Carregando modelo Whisper...")
    model = whisper.load_model("base")  # Pode mudar para 'small', 'medium', etc.

    print("Transcrevendo audio (mapeando palavras)...")
    result = model.transcribe(audio_path, word_timestamps=True)

    subtitles = []
    index = 1

    # 2. Extrai todas as palavras com seus respectivos tempos
    todas_palavras = []
    for segment in result["segments"]:
        if "words" in segment:
            todas_palavras.extend(segment["words"])

    if not todas_palavras:
        print("Nenhuma palavra com timestamp encontrada. Verifique o audio.")
        return

    # 3. Agrupa as palavras respeitando o limite de caracteres
    linha_atual = []
    comprimento_atual = 0
    tempo_inicio = None

    for i, p in enumerate(todas_palavras):
        palavra = p["word"].strip()
        # Tratamento para pontuações grudadas que o Whisper gera
        palavra_com_espaco = p["word"]

        # Se for a primeira palavra da legenda, define o tempo de inicio
        if not linha_atual:
            tempo_inicio = p["start"]

        # Verifica o tamanho se adicionarmos esta palavra
        # Se ja tiver palavras na linha, soma 1 por causa do espaco
        adicional = len(palavra) + (1 if linha_atual else 0)

        # Se estourar 20 caracteres E ja tiver alguma palavra na linha, fecha a legenda anterior
        if comprimento_atual + adicional > max_chars and linha_atual:
            # Tempo de fim é o fim da última palavra adicionada com sucesso
            tempo_fim = todas_palavras[i - 1]["end"]
            texto_legenda = "".join(linha_atual).strip()

            # Adiciona ao formato SRT
            subtitles.append(
                srt.Subtitle(
                    index=index,
                    start=timedelta(seconds=tempo_inicio),
                    end=timedelta(seconds=tempo_fim),
                    content=texto_legenda
                )
            )
            index += 1

            # Reseta para a nova linha com a palavra atual
            linha_atual = [palavra_com_espaco]
            comprimento_atual = len(palavra)
            tempo_inicio = p["start"]
        else:
            # Caso caiba ou seja a primeira palavra (mesmo que maior que 20)
            linha_atual.append(palavra_com_espaco)
            comprimento_atual += adicional

    # Adiciona a ultima linha restante, se houver
    if linha_atual:
        tempo_fim = todas_palavras[-1]["end"]
        subtitles.append(
            srt.Subtitle(
                index=index,
                start=timedelta(seconds=tempo_inicio),
                end=timedelta(seconds=tempo_fim),
                content="".join(linha_atual).strip()
            )
        )

    # 4. Grava o arquivo SRT final
    print(f"Salvando legenda em: {srt_path}")
    with open(srt_path, "w", encoding="utf-8") as f:
        f.write(srt.compose(subtitles))
    print("Processo concluido com sucesso!")


# --- COMO EXECUTAR ---
# Substitua pelos caminhos dos seus arquivos
criar_legendas_curtas("../proj_011/audio_original.mp3", "../proj_011/legenda_curta.srt", max_chars=17)