"""
Genera subtitulos animados (palabra por palabra, estilo TikTok/Reels) para un
video, usando los colores del sistema Taxes (resources/css/new_css_2025.css),
y los quema (hardcode) sobre el video con ffmpeg.

Requiere:
  - Python con el paquete `faster-whisper` instalado
  - ffmpeg/ffprobe disponibles (PATH o alguna ruta conocida de winget)

Uso:
  python generar_video_subtitulado.py video.mp4
  python generar_video_subtitulado.py video.mp4 --color success --palabras-por-linea 3
  python generar_video_subtitulado.py video.mp4 --no-quemar   # solo genera .srt/.ass para revisar

Flujo:
  1. Extrae el audio del video con ffmpeg (wav 16kHz mono).
  2. Transcribe con faster-whisper pidiendo timestamps por PALABRA.
  3. Guarda un .srt "plano" (por frase) para que el usuario pueda revisar y
     corregir errores de la transcripcion automatica antes de dar el video
     por bueno.
  4. Agrupa las palabras en lineas cortas y genera un .ass con animacion:
     las palabras van apareciendo a medida que se hablan, la palabra actual
     resaltada con un color del sistema y un pequeño "rebote" de escala.
  5. Quema el .ass sobre el video (salvo que se pida --no-quemar).
"""

import argparse
import json
import shutil
import subprocess
import sys
from pathlib import Path

# Paleta del sistema Taxes (resources/css/new_css_2025.css, bloque :root)
SYSTEM_COLORS = {
    "primary": "#2563eb",
    "success": "#10b981",
    "warning": "#f59e0b",
    "danger": "#ef4444",
    "info": "#06b6d4",
}
COLOR_WHITE_HEX = "#ffffff"
COLOR_OUTLINE_HEX = "#0f172a"  # --gray-900

# Rutas conocidas por si el binario no esta en PATH (instalacion via winget)
FFMPEG_FALLBACKS = [
    r"C:\Users\larza\AppData\Local\Microsoft\WinGet\Packages\Gyan.FFmpeg_Microsoft.Winget.Source_8wekyb3d8bbwe\ffmpeg-8.1.2-full_build\bin\ffmpeg.exe",
]
FFPROBE_FALLBACKS = [
    r"C:\Users\larza\AppData\Local\Microsoft\WinGet\Packages\Gyan.FFmpeg_Microsoft.Winget.Source_8wekyb3d8bbwe\ffmpeg-8.1.2-full_build\bin\ffprobe.exe",
]


def find_binary(name, fallbacks):
    found = shutil.which(name)
    if found:
        return found
    for candidate in fallbacks:
        if Path(candidate).exists():
            return candidate
    raise FileNotFoundError(
        f"No se encontro '{name}' en PATH ni en las rutas conocidas: {fallbacks}"
    )


def hex_to_ass_color(hex_color):
    """#RRGGBB -> &H00BBGGRR& (formato de color de libass/ASS)."""
    h = hex_color.lstrip("#")
    r, g, b = h[0:2], h[2:4], h[4:6]
    return f"&H00{b}{g}{r}&".upper()


def get_video_resolution(video_path, ffprobe_bin):
    out = subprocess.check_output(
        [
            ffprobe_bin, "-v", "error", "-select_streams", "v:0",
            "-show_entries", "stream=width,height", "-of", "csv=p=0",
            str(video_path),
        ],
        text=True,
    ).strip()
    w, h = out.split(",")
    return int(w), int(h)


def extract_audio(video_path, audio_path, ffmpeg_bin):
    subprocess.run(
        [
            ffmpeg_bin, "-y", "-i", str(video_path), "-vn",
            "-acodec", "pcm_s16le", "-ar", "16000", "-ac", "1", str(audio_path),
        ],
        check=True, capture_output=True,
    )


def transcribe_words(audio_path, model_size, language):
    from faster_whisper import WhisperModel

    model = WhisperModel(model_size, device="cpu", compute_type="int8")
    segments, info = model.transcribe(
        str(audio_path), language=language, vad_filter=True, word_timestamps=True
    )
    words = []
    for seg in segments:
        for w in seg.words:
            words.append({"word": w.word.strip(), "start": w.start, "end": w.end})
    return words, info


def to_srt_time(t):
    h = int(t // 3600)
    m = int((t % 3600) // 60)
    s = int(t % 60)
    ms = int(round((t - int(t)) * 1000))
    return f"{h:02d}:{m:02d}:{s:02d},{ms:03d}"


def to_ass_time(t):
    h = int(t // 3600)
    m = int((t % 3600) // 60)
    s = t % 60
    return f"{h:d}:{m:02d}:{s:05.2f}"


def group_words(words, max_words_per_line, pause_gap):
    groups, current = [], []
    for i, w in enumerate(words):
        current.append(w)
        ends_sentence = w["word"].rstrip().endswith((".", "!", "?"))
        next_gap = (words[i + 1]["start"] - w["end"]) if i + 1 < len(words) else None
        big_pause = next_gap is not None and next_gap > pause_gap
        if len(current) >= max_words_per_line or ends_sentence or big_pause or i == len(words) - 1:
            groups.append(current)
            current = []
    return groups


def write_srt(groups, srt_path):
    lines = []
    for idx, group in enumerate(groups, start=1):
        text = " ".join(w["word"] for w in group)
        lines.append(str(idx))
        lines.append(f"{to_srt_time(group[0]['start'])} --> {to_srt_time(group[-1]['end'])}")
        lines.append(text)
        lines.append("")
    srt_path.write_text("\n".join(lines), encoding="utf-8")


def write_ass(groups, ass_path, video_w, video_h, highlight_hex):
    color_white = hex_to_ass_color(COLOR_WHITE_HEX)
    color_highlight = hex_to_ass_color(highlight_hex)
    color_outline = hex_to_ass_color(COLOR_OUTLINE_HEX)

    fontsize = round(video_h * 0.065)
    margin_v = round(video_h * 0.12)

    header = f"""[Script Info]
ScriptType: v4.00+
PlayResX: {video_w}
PlayResY: {video_h}
ScaledBorderAndShadow: yes

[V4+ Styles]
Format: Name, Fontname, Fontsize, PrimaryColour, SecondaryColour, OutlineColour, BackColour, Bold, Italic, Underline, StrikeOut, ScaleX, ScaleY, Spacing, Angle, BorderStyle, Outline, Shadow, Alignment, MarginL, MarginR, MarginV, Encoding
Style: Default,Arial,{fontsize},{color_white},{color_white},{color_outline},&H00000000&,1,0,0,0,100,100,0,0,1,4,1,2,60,60,{margin_v},1

[Events]
Format: Layer, Start, End, Style, Name, MarginL, MarginR, MarginV, Effect, Text
"""

    events = []
    n_groups = len(groups)
    for gi, group in enumerate(groups):
        group_last_end = group[-1]["end"]
        next_group_start = groups[gi + 1][0]["start"] if gi + 1 < n_groups else group_last_end + 0.4
        natural_end = min(group_last_end + 0.35, next_group_start)

        for j in range(len(group)):
            start = group[j]["start"]
            end = group[j + 1]["start"] if j + 1 < len(group) else natural_end
            if end <= start:
                end = start + 0.05

            pieces = []
            for k in range(j + 1):
                word_text = group[k]["word"]
                if k == j:
                    pieces.append(
                        "{\\c" + color_highlight
                        + "\\fscx55\\fscy55\\t(0,110,\\fscx112\\fscy112)\\t(110,190,\\fscx100\\fscy100)}"
                        + word_text + "{\\r}"
                    )
                else:
                    pieces.append("{\\c" + color_white + "}" + word_text)
            text = " ".join(pieces)
            events.append(f"Dialogue: 0,{to_ass_time(start)},{to_ass_time(end)},Default,,0,0,0,,{text}")

    ass_path.write_text(header + "\n".join(events) + "\n", encoding="utf-8")


def burn_subtitles(video_path, ass_path, output_path, ffmpeg_bin):
    # Se ejecuta con cwd = carpeta del ass para evitar problemas de escapeo
    # de rutas (":" y "\") del filtro "ass" de ffmpeg en Windows.
    subprocess.run(
        [
            ffmpeg_bin, "-y", "-i", str(video_path.resolve()),
            "-vf", f"ass={ass_path.name}",
            "-c:a", "copy", str(output_path.resolve()),
        ],
        check=True, capture_output=True, cwd=str(ass_path.parent),
    )


def main():
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("video", type=Path, help="Video de entrada (mp4)")
    parser.add_argument("--modelo", default="small", help="Modelo whisper (tiny/base/small/medium). Default: small")
    parser.add_argument("--idioma", default="es", help="Idioma del audio. Default: es")
    parser.add_argument("--color", default="primary", choices=SYSTEM_COLORS.keys(),
                         help="Color del sistema para resaltar la palabra actual. Default: primary")
    parser.add_argument("--palabras-por-linea", type=int, default=4, dest="max_words")
    parser.add_argument("--pausa-corte", type=float, default=0.45, dest="pause_gap",
                         help="Segundos de silencio que fuerzan un corte de linea")
    parser.add_argument("--no-quemar", action="store_true", dest="no_burn",
                         help="Solo genera .srt/.ass para revisar, no quema el video")
    parser.add_argument("--mantener-audio", action="store_true", dest="keep_audio")
    args = parser.parse_args()

    video_path = args.video
    if not video_path.exists():
        sys.exit(f"No existe el video: {video_path}")

    out_dir = video_path.parent
    stem = video_path.stem
    audio_path = out_dir / f"{stem}_audio.wav"
    srt_path = out_dir / f"{stem}_subtitulos.srt"
    ass_path = out_dir / f"{stem}_subtitulos_animados.ass"
    output_video_path = out_dir / f"{stem}_subtitulado.mp4"

    ffmpeg_bin = find_binary("ffmpeg", FFMPEG_FALLBACKS)
    ffprobe_bin = find_binary("ffprobe", FFPROBE_FALLBACKS)

    print("1/5 Extrayendo audio...")
    extract_audio(video_path, audio_path, ffmpeg_bin)

    print(f"2/5 Transcribiendo con faster-whisper ({args.modelo}, {args.idioma})...")
    words, info = transcribe_words(audio_path, args.modelo, args.idioma)
    print(f"    {len(words)} palabras detectadas (idioma: {info.language}, prob {info.language_probability:.2f})")

    if not args.keep_audio:
        audio_path.unlink(missing_ok=True)

    groups = group_words(words, args.max_words, args.pause_gap)

    print(f"3/5 Guardando transcripcion en {srt_path.name} (revisar antes de compartir)")
    write_srt(groups, srt_path)

    print(f"4/5 Generando subtitulos animados en {ass_path.name} (color: {args.color})")
    video_w, video_h = get_video_resolution(video_path, ffprobe_bin)
    write_ass(groups, ass_path, video_w, video_h, SYSTEM_COLORS[args.color])

    if args.no_burn:
        print("5/5 --no-quemar: listo, revisa el .srt y volve a correr sin ese flag para generar el video final.")
        return

    print(f"5/5 Quemando subtitulos sobre el video -> {output_video_path.name}")
    burn_subtitles(video_path, ass_path, output_video_path, ffmpeg_bin)
    print("Listo:", output_video_path)


if __name__ == "__main__":
    main()