subtitle_rendering.py

from __future__ import annotations

import math
import uuid

import moviepy.editor as mpe
import numpy as np
from PIL import Image, ImageDraw, ImageFont, ImageColor

# ----------------------------
# Шрыфты
# ----------------------------
AVAILABLE_FONTS = {
    "DejaVuSans-Bold": "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf",
    "DejaVuSans": "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf",
    "LiberationSerif-Bold": "/usr/share/fonts/truetype/liberation/LiberationSerif-Bold.ttf",
    "BadScript-Regular": "fonts/Bad_Script/BadScript-Regular.ttf",
    "Gidole-Regular": "fonts/Gidole/Gidole-Regular.ttf",
    "GreatVibes-Regular": "fonts/Great_Vibes/GreatVibes-Regular.ttf",
    "OpenSans-Variable": "fonts/Open_Sans/OpenSans-VariableFont_wdth,wght.ttf",
    "OpenSans-Italic-Variable": "fonts/Open_Sans/OpenSans-Italic-VariableFont_wdth,wght.ttf",
    "Roboto-Variable": "fonts/Roboto/Roboto-VariableFont_wdth,wght.ttf",
    "Roboto-Italic-Variable": "fonts/Roboto/Roboto-Italic-VariableFont_wdth,wght.ttf",
    "SourceCodePro-Variable": "fonts/Source_Code_Pro/SourceCodePro-VariableFont_wght.ttf",
    "SourceCodePro-Italic-Variable": "fonts/Source_Code_Pro/SourceCodePro-Italic-VariableFont_wght.ttf",
    "Tektur-Variable": "fonts/Tektur/Tektur-VariableFont_wdth,wght.ttf",
    "Ponomar-Regular": "fonts/Ponomar/Ponomar-Regular.ttf",
}


# ----------------------------
# Колер -> RGB
# ----------------------------
def hex_to_rgb(color_str: str):
    try:
        return ImageColor.getrgb(color_str)
    except Exception:
        pass

    lower = (color_str or "").lower()
    if lower.startswith("rgba") or lower.startswith("rgb"):
        inside = color_str[color_str.find("(") + 1 : color_str.rfind(")")]
        parts = [p.strip() for p in inside.split(",")]
        if len(parts) >= 3:
            try:
                r, g, b = [int(float(parts[i])) for i in range(3)]
                return (r, g, b)
            except Exception:
                pass

    c = (color_str or "#FFFFFF").lstrip("#")
    if len(c) == 3:
        c = "".join([ch * 2 for ch in c])
    return tuple(int(c[i : i + 2], 16) for i in (0, 2, 4))


# ----------------------------
# SRT-парсінг
# ----------------------------
def srt_time_to_sec(time_str):
    t = time_str.strip().replace(".", ",")
    if "," not in t:
        t += ",000"
    h, m, s_ms = t.split(":")
    s, ms = s_ms.split(",")
    return int(h) * 3600 + int(m) * 60 + int(s) + int(ms) / 1000


def parse_srt_from_text(srt_text):
    subs = []
    normalized = (srt_text or "").replace("\r\n", "\n").replace("\r", "\n")
    for block in normalized.strip().split("\n\n"):
        lines = block.splitlines()
        if len(lines) < 3:
            continue
        times = lines[1]
        try:
            start_str, end_str = times.split("-->")
            start = srt_time_to_sec(start_str)
            end = srt_time_to_sec(end_str)
            text = "\n".join(lines[2:]).strip()  # захоўваем пераносы
            subs.append({"start": start, "end": end, "text": text})
        except Exception:
            continue
    return subs


# ----------------------------
# Wrap: аўтаматычны перанос радкоў па шырыні кадра
# ----------------------------
def _line_width_px(draw: ImageDraw.ImageDraw, line: str, font: ImageFont.ImageFont, stroke_width: int = 0) -> int:
    try:
        l, t, r, b = draw.textbbox((0, 0), line, font=font, stroke_width=stroke_width)
        return int(r - l)
    except Exception:
        try:
            w = draw.textlength(line, font=font)
            return int(w + 2 * stroke_width)
        except Exception:
            w, _ = draw.textsize(line, font=font)
            return int(w + 2 * stroke_width)


def _break_long_word(draw, word: str, font, max_width: int, stroke_width: int) -> list:
    chunks = []
    cur = ""
    for ch in word:
        test = cur + ch
        if cur and _line_width_px(draw, test, font, stroke_width) > max_width:
            chunks.append(cur)
            cur = ch
        else:
            cur = test
    if cur:
        chunks.append(cur)
    return chunks


def wrap_text_to_width(text: str, draw: ImageDraw.ImageDraw, font: ImageFont.ImageFont,
                       max_width: int, stroke_width: int = 0) -> str:
    if not text:
        return text

    paragraphs = (text or "").split("\n")
    out_lines = []

    for para in paragraphs:
        p = para.strip()
        if not p:
            out_lines.append("")
            continue

        words = p.split()
        line = ""
        for w in words:
            if not line:
                if _line_width_px(draw, w, font, stroke_width) <= max_width:
                    line = w
                else:
                    chunks = _break_long_word(draw, w, font, max_width, stroke_width)
                    out_lines.extend(chunks[:-1])
                    line = chunks[-1] if chunks else ""
                continue

            test = f"{line} {w}"
            if _line_width_px(draw, test, font, stroke_width) <= max_width:
                line = test
            else:
                out_lines.append(line)
                if _line_width_px(draw, w, font, stroke_width) <= max_width:
                    line = w
                else:
                    chunks = _break_long_word(draw, w, font, max_width, stroke_width)
                    out_lines.extend(chunks[:-1])
                    line = chunks[-1] if chunks else ""

        if line:
            out_lines.append(line)

    return "\n".join(out_lines)


# ----------------------------
# Тэкставы кліп (PIL -> moviepy), без абразання
# ----------------------------
def _measure_multiline_bbox(draw, text, font, stroke_width=0, spacing=4):
    try:
        return draw.multiline_textbbox(
            (0, 0),
            text,
            font=font,
            stroke_width=stroke_width,
            spacing=spacing,
            align="center",
        )
    except Exception:
        pass

    try:
        return draw.textbbox((0, 0), text, font=font, stroke_width=stroke_width)
    except Exception:
        pass

    w, h = draw.textsize(text, font=font)
    h = h + max(2, stroke_width + 2)
    return (0, 0, w, h)


def _clamp(v: float, lo: float, hi: float) -> float:
    return max(lo, min(hi, v))


def create_animated_text_clip(
    text,
    duration,
    font,
    fontsize,
    color,
    stroke_color,
    stroke_width,
    position_type,
    custom_x_shift,  # ЗРУХ X (адносна цэнтра)
    custom_y,
    animation,
    video_width,
    video_height,
    bg_color=None,
    bg_opacity=1.0,
    wrap_ratio: float = 0.90,
):
    try:
        txt_rgb = hex_to_rgb(color)
    except Exception:
        txt_rgb = (255, 255, 255)

    try:
        stroke_rgb = hex_to_rgb(stroke_color)
    except Exception:
        stroke_rgb = (0, 0, 0)

    if bg_color:
        try:
            bg_rgb = hex_to_rgb(bg_color)
        except Exception:
            bg_rgb = (0, 0, 0)
        bg_a = int(max(0.0, min(1.0, bg_opacity)) * 255)
    else:
        bg_rgb = (0, 0, 0)
        bg_a = 0

    try:
        pil_font = ImageFont.truetype(font, fontsize)
    except Exception:
        pil_font = ImageFont.load_default()

    dummy_img = Image.new("RGBA", (10, 10), (0, 0, 0, 0))
    dummy_draw = ImageDraw.Draw(dummy_img)

    spacing = max(0, int(fontsize * 0.15))

    # --- AUTO WRAP ---
    max_text_width = max(200, int(video_width * float(wrap_ratio)))
    wrapped_text = wrap_text_to_width(text, dummy_draw, pil_font, max_text_width, stroke_width=stroke_width)

    # --- BBOX ---
    l, t, r, b = _measure_multiline_bbox(dummy_draw, wrapped_text, pil_font, stroke_width=stroke_width, spacing=spacing)
    text_w = r - l
    text_h = b - t

    pad_x = max(10, int(fontsize * 0.25))
    pad_y = max(6, int(fontsize * 0.20))
    safe = max(2, int(stroke_width) + 3)

    img_w = int(math.ceil(text_w + 2 * (pad_x + safe)))
    img_h = int(math.ceil(text_h + 2 * (pad_y + safe)))

    img = Image.new("RGBA", (img_w, img_h), bg_rgb + (bg_a,))
    draw = ImageDraw.Draw(img)

    # галоўнае: улічваем bbox (можа быць адмоўны)
    text_x = (pad_x + safe) - l
    text_y = (pad_y + safe) - t

    draw.multiline_text(
        (text_x, text_y),
        wrapped_text,
        font=pil_font,
        fill=txt_rgb + (255,),
        stroke_width=stroke_width,
        stroke_fill=stroke_rgb + (255,),
        spacing=spacing,
        align="center",
    )

    arr = np.array(img)
    rgb = arr[..., :3]
    alpha = arr[..., 3].astype(np.float32) / 255.0

    clip = mpe.ImageClip(rgb, ismask=False).set_duration(duration)
    mask = mpe.ImageClip(alpha, ismask=True).set_duration(duration)
    clip = clip.set_mask(mask)

    # --- Пазіцыя (аптымізавана) ---
    # верх/ніз ссоўваем да цэнтра на 15% вышыні кадра
    center_pull = 0.15 * float(video_height)
    base_x_center = (video_width - clip.w) / 2

    if position_type == "bottom":
        x = base_x_center
        y = (video_height - clip.h - 20) - center_pull
    elif position_type == "top":
        x = base_x_center
        y = 20 + center_pull
    elif position_type == "center":
        x = base_x_center
        y = (video_height - clip.h) / 2
    else:  # custom: Х заўсёды ад цэнтра (як ва ўсіх), custom_x_shift — дадатковы зрух
        x = base_x_center + float(custom_x_shift)
        y = float(custom_y)

    # clamp, каб не вылятала за кадр
    x = _clamp(x, 0, max(0, video_width - clip.w))
    y = _clamp(y, 0, max(0, video_height - clip.h))

    anim = (animation or "").lower()

    if anim == "fade":
        fd = min(0.5, duration / 2)
        clip = clip.fadein(fd).fadeout(fd)
        return clip.set_position((x, y))
    elif anim == "slide":
        fd = min(0.5, duration / 2)

        def slide_pos(t_):
            progress = min(max(t_ / fd, 0), 1)
            return -clip.w + (x + clip.w) * progress, y

        return clip.set_position(slide_pos)
    elif anim == "zoom":
        clip = clip.resize(lambda t_: 0.5 + 0.5 * min(max(t_ / duration, 0), 1))
        return clip.set_position((x, y))
    else:
        return clip.set_position((x, y))


# ----------------------------
# Накладанне субтытраў
# ----------------------------
def apply_subtitles(
    video_path,
    subtitles_text,
    font,
    fontsize,
    color,
    stroke_color,
    stroke_width,
    position_type,
    custom_x_shift,
    custom_y,
    animation,
    export_quality,
    bg_color=None,
    bg_opacity=1.0,
    wrap_ratio: float = 0.90,
):
    subs = parse_srt_from_text(subtitles_text)
    video = mpe.VideoFileClip(video_path)
    w, h = video.w, video.h

    clips = [video]
    for s in subs:
        dur = s["end"] - s["start"]
        if dur <= 0:
            continue
        txt_clip = create_animated_text_clip(
            s["text"],
            dur,
            font,
            fontsize,
            color,
            stroke_color,
            stroke_width,
            position_type,
            custom_x_shift,
            custom_y,
            animation,
            w,
            h,
            bg_color=bg_color,
            bg_opacity=bg_opacity,
            wrap_ratio=wrap_ratio,
        ).set_start(s["start"])
        clips.append(txt_clip)

    final = mpe.CompositeVideoClip(clips)

    if export_quality == "мінімальнае":
        final = final.resize(height=480)
    elif export_quality == "сярэдняе":
        final = final.resize(height=720)
    elif export_quality == "максімальнае" and h < 1080:
        final = final.resize(height=1080)

    out = f"output_video_{uuid.uuid4().hex}.mp4"
    try:
        final.write_videofile(out, fps=video.fps, codec="libx264", audio_codec="aac")
    finally:
        video.close()
        final.close()
    return out


# ----------------------------
# Перадпрагляд (1 секунда)
# ----------------------------
def extract_first_frame(video_path):
    try:
        clip = mpe.VideoFileClip(video_path)
        frame = clip.get_frame(0)
        clip.close()
        return frame
    except Exception:
        return None


def create_single_frame_video(
    frame,
    text,
    font,
    fontsize,
    color,
    stroke_color,
    stroke_width,
    position_type,
    custom_x_shift,
    custom_y,
    animation,
    bg_color=None,
    bg_opacity=1.0,
    wrap_ratio: float = 0.90,
):
    if frame is None:
        return None

    h, w, _ = frame.shape
    base = mpe.ImageClip(frame).set_duration(1.0)

    txt_clip = create_animated_text_clip(
        text,
        1.0,
        font,
        fontsize,
        color,
        stroke_color,
        stroke_width,
        position_type,
        custom_x_shift,
        custom_y,
        animation,
        w,
        h,
        bg_color=bg_color,
        bg_opacity=bg_opacity,
        wrap_ratio=wrap_ratio,
    )

    final_clip = mpe.CompositeVideoClip([base, txt_clip])
    path = f"preview_video_{uuid.uuid4().hex}.mp4"
    try:
        final_clip.write_videofile(
            path,
            fps=24,
            codec="libx264",
            audio=False,
            verbose=False,
            logger=None,
        )
    finally:
        final_clip.close()
    return path