#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import re
from pathlib import Path

TARGET = Path("manga-translator.py")

def cut_after_first_entrypoint(text: str) -> str:
    """
    Keep only first full __main__ block and remove duplicated tail if present.
    """
    m = re.search(r'(?m)^if __name__ == "__main__":\s*$', text)
    if not m:
        return text

    start = m.start()
    # Keep entrypoint block plus indented lines after it
    lines = text[start:].splitlines(True)
    keep = []
    keep.append(lines[0])  # if __name__...
    i = 1
    while i < len(lines):
        ln = lines[i]
        if ln.strip() == "":
            keep.append(ln)
            i += 1
            continue
        # if dedented back to col 0 => end of block
        if not ln.startswith((" ", "\t")):
            break
        keep.append(ln)
        i += 1

    cleaned = text[:start] + "".join(keep)
    return cleaned

def replace_bad_vars(text: str) -> str:
    text = text.replace(
        "merge_micro_boxes_relaxed(bubbles, bubble_boxes, bubble_quads, bubble_indices, ocr, image_bgr)",
        "merge_micro_boxes_relaxed(bubbles, bubble_boxes, bubble_quads, bubble_indices, filtered, image)"
    )
    text = text.replace(
        "reattach_orphan_short_tokens(bubbles, bubble_boxes, bubble_quads, bubble_indices, ocr)",
        "reattach_orphan_short_tokens(bubbles, bubble_boxes, bubble_quads, bubble_indices, filtered)"
    )
    return text

def ensure_autofix_chain(text: str) -> str:
    old = (
        "    # ── Auto-fix (split + merge) ──────────────────────────────────────────\n"
        "    if auto_fix_bubbles:\n"
        "        bubbles, bubble_boxes, bubble_quads, bubble_indices = merge_micro_boxes_relaxed(bubbles, bubble_boxes, bubble_quads, bubble_indices, filtered, image)\n"
    )
    new = (
        "    # ── Auto-fix (split + merge) ──────────────────────────────────────────\n"
        "    if auto_fix_bubbles:\n"
        "        bubbles, bubble_boxes, bubble_quads, bubble_indices = auto_fix_bubble_detection(\n"
        "            bubble_boxes, bubble_indices, bubble_quads, bubbles, filtered, image)\n"
        "        bubbles, bubble_boxes, bubble_quads, bubble_indices = merge_micro_boxes_relaxed(\n"
        "            bubbles, bubble_boxes, bubble_quads, bubble_indices, filtered, image)\n"
    )
    return text.replace(old, new)

def ensure_split_commit(text: str) -> str:
    marker = "    # ── Remove nested / duplicate boxes ──────────────────────────────────\n"
    if marker not in text:
        return text

    if "bubbles = new_bubbles" in text:
        return text

    inject = (
        "    bubbles = new_bubbles\n"
        "    bubble_boxes = new_bubble_boxes\n"
        "    bubble_quads = new_bubble_quads\n"
        "    bubble_indices = new_bubble_indices\n\n"
    )
    return text.replace(marker, inject + marker)

def ensure_rescue_pipeline(text: str) -> str:
    anchor = '    print(f"Kept: {len(filtered)} | Skipped: {skipped}")\n'
    if anchor not in text:
        return text

    if "rescue_name_and_short_tokens(raw" in text:
        return text

    block = (
        '    print(f"Kept: {len(filtered)} | Skipped: {skipped}")\n'
        '    # Protect short dialogue tokens confidence\n'
        '    tmp = []\n'
        '    for bbox, t, conf in filtered:\n'
        '        tmp.append((bbox, t, maybe_conf_floor_for_protected(t, conf, floor=0.40)))\n'
        '    filtered = tmp\n'
        '    # Rescue names/short tokens dropped by strict filters\n'
        '    rescued = rescue_name_and_short_tokens(raw, min_conf=0.20)\n'
        '    filtered = merge_rescued_items(filtered, rescued, iou_threshold=0.55)\n'
    )
    return text.replace(anchor, block)

def main():
    if not TARGET.exists():
        raise FileNotFoundError(f"Not found: {TARGET}")

    src = TARGET.read_text(encoding="utf-8")
    out = src

    out = cut_after_first_entrypoint(out)
    out = replace_bad_vars(out)
    out = ensure_autofix_chain(out)
    out = ensure_split_commit(out)
    out = ensure_rescue_pipeline(out)

    TARGET.write_text(out, encoding="utf-8")
    print("✅ Patched manga-translator.py")

if __name__ == "__main__":
    main()