ui.py

from __future__ import annotations

import os

import gradio as gr

from subtitle_generation import process_video_to_srt, read_srt_file
from subtitle_rendering import (
    AVAILABLE_FONTS,
    apply_subtitles,
    create_single_frame_video,
    extract_first_frame,
    parse_srt_from_text,
)

PLAYER_H = 520
CSS = f"""
#main_video, #main_video > div {{
  height: {PLAYER_H}px !important;
}}
#main_video video {{
  height: {PLAYER_H}px !important;
  width: 100% !important;
  object-fit: contain !important;
  background: #000 !important;
}}
"""


def build_ui():
    with gr.Blocks(css=CSS) as demo:
        gr.Markdown("# Аўтаматычнае стварэнне і накладанне субтытраў")
    
        # Стан: арыгінальны файл (каб не губляць пры preview/output)
        uploaded_video = gr.State(None)
        first_frame = gr.State(None)
    
        with gr.Row(equal_height=True):
            with gr.Column(scale=6):
                video_input = gr.Video(
                    label="Відэа (тут жа будзе перадпрагляд і вынік)",
                    elem_id="main_video",
                )
                with gr.Row():
                    show_original_btn = gr.Button("Паказаць арыгінал")
                    # (кнопка дадатковая, але карысная пасля preview)
    
            with gr.Column(scale=5):
                # 1) Стварэнне субтытраў — разварочваецца/згортваецца
                with gr.Accordion("Стварэнне субтытраў", open=False):
                    length_dd = gr.Dropdown(
                        ["Кароткія (1 слова)", "Сярэднія (2-3 словы)", "Доўгія (сказы)"],
                        value="Сярэднія (2-3 словы)",
                        label="Даўжыня субцітру",
                    )
                    srt_file = gr.File(
                        label="Імпартаваць гатовы SRT",
                        file_types=[".srt"],
                        type="filepath",
                    )
                    gen_btn = gr.Button("Стварыць субтытры")
                    subs_box = gr.Textbox(label="Тэкст субтытраў", lines=14)
    
                # 2) Стыль
                with gr.Accordion("Стыль субтытраў", open=True):
                    with gr.Row():
                        font_dd = gr.Dropdown(list(AVAILABLE_FONTS.keys()), value="DejaVuSans-Bold", label="Шрыфт")
                        size_sl = gr.Slider(10, 100, 1, value=40, label="Памер")
    
                    with gr.Row():
                        color_txt = gr.ColorPicker(value="#FFFFFF", label="Колер тэксту")
                        color_st = gr.ColorPicker(value="#000000", label="Колер абводкі")
                        width_sl = gr.Slider(0, 10, 1, value=2, label="Абводка")
    
                    with gr.Row():
                        bg_color = gr.ColorPicker(value="#000000", label="Фон")
                        bg_opac = gr.Slider(0.0, 1.0, 0.1, value=0.5, label="Празрыстасць")
    
                    with gr.Row():
                        pos_dd = gr.Dropdown(["bottom", "top", "center", "custom"], value="bottom", label="Пазіцыя")
                        anim_dd = gr.Dropdown(["None", "fade", "slide", "zoom"], value="None", label="Анімацыя")
    
                    # Custom controls (паказваем толькі калі патрэбна)
                    with gr.Row():
                        x_shift_sl = gr.Slider(-960, 960, 1, value=0, label="X-зрух ад цэнтра", visible=False)
                        y_sl = gr.Slider(0, 1080, 1, value=100, label="Y (для custom)", visible=False)
    
                    with gr.Row():
                        wrap_en = gr.Checkbox(value=True, label="Аўта-перанос радкоў")
                        wrap_pct = gr.Slider(50, 100, 1, value=90, label="Макс. шырыня радка (%)")
    
                # 3) Экспарт
                with gr.Accordion("Экспарт", open=False):
                    exp_dd = gr.Dropdown(
                        ["мінімальнае", "арыгінальнае", "сярэдняе", "максімальнае"],
                        value="арыгінальнае",
                        label="Якасць",
                    )
                    apply_btn = gr.Button("Накласці субтытры (фінальны вынік)")
    
        # --- Лагічныя апрацоўшчыкі ---
    
        # Калі значэнне video_input змянілася:
        # - калі гэта preview/output -> НЕ пераціраем uploaded_video
        # - калі гэта арыгінал (загрузка) -> абнаўляем uploaded_video і first_frame
        def _on_video_change(v, current_uploaded, current_frame):
            if not v:
                return None, None
    
            # калі проста вярнуліся да арыгінала праз кнопку і ён ужо вядомы — не пералічваем кадр
            if current_uploaded and str(v) == str(current_uploaded):
                return current_uploaded, current_frame
    
            name = os.path.basename(str(v))
            if name.startswith(("preview_video_", "output_video_")):
                return current_uploaded, current_frame
    
            frame = extract_first_frame(v)
            return v, frame
    
        video_input.change(
            _on_video_change,
            inputs=[video_input, uploaded_video, first_frame],
            outputs=[uploaded_video, first_frame],
        )
    
        # Паказаць арыгінал у тым жа плэеры
        def _show_original(v_orig):
            return v_orig
    
        show_original_btn.click(_show_original, inputs=uploaded_video, outputs=video_input)
    
        # Генерацыя субтытраў заўсёды па арыгінальным uploaded_video
        def _map_length_choice(label):
            if label == "Кароткія (1 слова)":
                return "short"
            if label == "Доўгія (сказы)":
                return "long"
            return "medium"

        def _on_generate_click(v_orig, length_label):
            if not v_orig:
                return "Не выбрана відэа"
            result, _ = process_video_to_srt(v_orig, length_mode=_map_length_choice(length_label))
            return result

        gen_btn.click(_on_generate_click, inputs=[uploaded_video, length_dd], outputs=subs_box)

        def _on_srt_upload(path):
            if not path:
                return ""
            return read_srt_file(path)

        srt_file.change(_on_srt_upload, inputs=srt_file, outputs=subs_box)
    
        # Пераключэнне бачнасці custom-слайдэраў
        def _toggle_custom(pos):
            vis = (pos == "custom")
            return gr.update(visible=vis), gr.update(visible=vis)
    
        pos_dd.change(_toggle_custom, inputs=pos_dd, outputs=[x_shift_sl, y_sl])
    
        # Перадпрагляд: 1 секунда на першым кадры, у тым жа video_input
        def update_preview(frame, text, font, size, col, colst, wd, pos, x_shift, y, anim, bgc, bgo, do_wrap, w_pct):
            if frame is None:
                return None
    
            subs = parse_srt_from_text(text)
            txt = subs[0]["text"] if subs else "Preview"
            fp = AVAILABLE_FONTS.get(font, font)
    
            wrap_ratio = float(w_pct) / 100.0 if do_wrap else 10.0  # вельмі вялікае -> фактычна без wrap
    
            return create_single_frame_video(
                frame,
                txt,
                fp,
                size,
                col,
                colst,
                wd,
                pos,
                x_shift,
                y,
                anim,
                bg_color=bgc,
                bg_opacity=bgo,
                wrap_ratio=wrap_ratio,
            )
    
        preview_inputs = [
            first_frame,
            subs_box,
            font_dd,
            size_sl,
            color_txt,
            color_st,
            width_sl,
            pos_dd,
            x_shift_sl,
            y_sl,
            anim_dd,
            bg_color,
            bg_opac,
            wrap_en,
            wrap_pct,
        ]
    
        # Любая змена ў стылі/тэксце абнаўляе перадпрагляд
        for c in [subs_box, font_dd, size_sl, color_txt, color_st, width_sl, pos_dd, x_shift_sl, y_sl, anim_dd, bg_color, bg_opac, wrap_en, wrap_pct]:
            c.change(update_preview, inputs=preview_inputs, outputs=video_input)
    
        # Фінальны рэндэр у тым жа плэеры, але па арыгінальным uploaded_video
        def _on_apply(
            v_orig,
            t,
            f,
            s,
            ct,
            cs,
            wd,
            p,
            x_shift,
            y,
            a,
            eq,
            bgc,
            bgo,
            do_wrap,
            w_pct,
        ):
            if not v_orig:
                return None
            font_path = AVAILABLE_FONTS.get(f, f)
            wrap_ratio = float(w_pct) / 100.0 if do_wrap else 10.0
            return apply_subtitles(
                v_orig,
                t,
                font_path,
                s,
                ct,
                cs,
                wd,
                p,
                x_shift,
                y,
                a,
                eq,
                bg_color=bgc,
                bg_opacity=bgo,
                wrap_ratio=wrap_ratio,
            )
    
        apply_btn.click(
            _on_apply,
            inputs=[
                uploaded_video,
                subs_box,
                font_dd,
                size_sl,
                color_txt,
                color_st,
                width_sl,
                pos_dd,
                x_shift_sl,
                y_sl,
                anim_dd,
                exp_dd,
                bg_color,
                bg_opac,
                wrap_en,
                wrap_pct,
            ],
            outputs=video_input,
        )
    
        return demo