import cv2
import numpy as np
from PIL import Image, ImageDraw, ImageFont

# ============================================================
# 字幕解析
# ============================================================

def parse_srt(srt_text: str):
    segments = []
    for block in srt_text.strip().split("\n\n"):
        lines = block.strip().split("\n")
        if len(lines) >= 3:
            times = lines[1].split(" --> ")
            segments.append((timestamp_to_seconds(times[0]), timestamp_to_seconds(times[1]), "\n".join(lines[2:])))
    return segments

def timestamp_to_seconds(ts: str):
    ts = ts.replace(",", ".")
    h, m, s = ts.split(":")
    return int(h) * 3600 + int(m) * 60 + float(s)

# ============================================================
# 字幕绘制（白字黑边，固定样式，无特效参数）
# ============================================================

def _draw_subtitle_lines(frame, text, font_path, video_w, video_h, font_size=None):
    """内部通用：白色文字 + 黑色描边，居中底部显示"""
    if not text:
        return frame

    img_rgba = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)).convert("RGBA")
    draw = ImageDraw.Draw(img_rgba)

    if font_size is None or font_size <= 0:
        font_size = max(18, min(video_w // 22, 40))

    try:
        font_obj = ImageFont.truetype(font_path, font_size) if font_path else ImageFont.load_default()
    except:
        font_obj = ImageFont.load_default()

    # 自动换行
    max_text_width = int(video_w * 0.85)
    lines, current_line = [], ""
    for char in text:
        test_line = current_line + char
        if draw.textbbox((0, 0), test_line, font=font_obj)[2] > max_text_width and current_line:
            lines.append(current_line)
            current_line = char
        else:
            current_line = test_line
    if current_line:
        lines.append(current_line)

    line_height = font_size + 10
    bottom_margin = int(video_h * 0.08)
    y_start = video_h - bottom_margin - len(lines) * line_height

    for i, line in enumerate(lines):
        tw = draw.textbbox((0, 0), line, font=font_obj)[2]
        tx = (video_w - tw) // 2
        ty = y_start + i * line_height
        # 黑色描边
        draw.text((tx, ty), line, fill=(255, 255, 255, 255), font=font_obj,
                  stroke_width=2, stroke_fill=(0, 0, 0, 220))

    return cv2.cvtColor(np.array(img_rgba.convert("RGB")), cv2.COLOR_RGB2BGR)


def draw_subtitle(frame, text, font_path, video_w, video_h,
                  font_size=None, **kwargs):
    return _draw_subtitle_lines(frame, text, font_path, video_w, video_h, font_size)


def draw_typewriter_subtitle(frame, full_text, visible_chars, font_path, video_w, video_h,
                              font_size=None, **kwargs):
    if not full_text or visible_chars <= 0:
        return frame
    return _draw_subtitle_lines(frame, full_text[:visible_chars], font_path, video_w, video_h, font_size)