with open("backend/app/mkv.py", "r", encoding="utf-8") as f:
    content = f.read()

import re

old_sig = """def build_episodes(
    video_path: str,
    total_duration: float,
    intro_seconds: float,
    outro_seconds: float,
    cut_points: list[float],
    output_dir: Path,
    temp_dir: Path,
    project_dir: Path,
    output_prefix: str = "episode",
    reencode: bool = False,
    ffmpeg_pass1_template: str | None = None,
    ffmpeg_pass2_template: str | None = None,
    progress_cb: ProgressCallback | None = None,
    log_cb: LogCallback | None = None,
) -> list[str]:"""

new_sig = """def build_episodes(
    video_path: str,
    total_duration: float,
    intro_seconds: float,
    outro_seconds: float,
    cut_points: list[float],
    output_dir: Path,
    temp_dir: Path,
    project_dir: Path,
    output_prefix: str = "episode",
    reencode: bool = False,
    ffmpeg_pass1_template: str | None = None,
    ffmpeg_pass2_template: str | None = None,
    progress_cb: ProgressCallback | None = None,
    log_cb: LogCallback | None = None,
    custom_segments: list[tuple[float, float]] | None = None,
) -> list[str]:"""

content = content.replace(old_sig, new_sig)

old_logic = """    core_end = max(intro_seconds, total_duration - outro_seconds)
    boundaries = [p for p in sorted(cut_points) if intro_seconds < p < core_end]
    boundaries.append(core_end)

    outputs: list[str] = []
    prev = intro_seconds
    safe_boundaries: list[float] = []

    for end in boundaries:
        if end - prev <= min_segment:
            continue
        safe_boundaries.append(end)
        prev = end

    if not safe_boundaries:
        raise RuntimeError("No valid segments after filtering short ranges")

    prev = intro_seconds
    total_segments = len(safe_boundaries)
    core_ranges: list[tuple[int, float, float]] = []
    for index, end in enumerate(safe_boundaries, start=1):
        core_ranges.append((index, prev, end))
        prev = end"""

new_logic = """    outputs: list[str] = []
    core_ranges: list[tuple[int, float, float]] = []
    min_segment = 0.001

    if custom_segments and len(custom_segments) > 0:
        total_segments = len(custom_segments)
        for index, (start, end) in enumerate(custom_segments, start=1):
            core_ranges.append((index, start, end))
    else:
        core_end = max(intro_seconds, total_duration - outro_seconds)
        boundaries = [p for p in sorted(cut_points) if intro_seconds < p < core_end]
        boundaries.append(core_end)

        prev = intro_seconds
        safe_boundaries: list[float] = []

        for end in boundaries:
            if end - prev <= min_segment:
                continue
            safe_boundaries.append(end)
            prev = end

        if not safe_boundaries:
            raise RuntimeError("No valid segments after filtering short ranges")

        prev = intro_seconds
        total_segments = len(safe_boundaries)
        for index, end in enumerate(safe_boundaries, start=1):
            core_ranges.append((index, prev, end))
            prev = end"""

content = content.replace(old_logic, new_logic)

with open("backend/app/mkv.py", "w", encoding="utf-8") as f:
    f.write(content)
print("mkv.py patched")