#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import cv2
import numpy as np
from PIL import Image, ImageDraw, ImageFont
import textwrap
import os
import re
import json
import warnings
from typing import List, Tuple, Dict, Any, Optional

from deep_translator import GoogleTranslator

# macOS Native Vision imports
import Vision
import Quartz
from Foundation import NSData

warnings.filterwarnings("ignore", category=UserWarning)

# ============================================================
# CONFIG
# ============================================================
TOP_BAND_RATIO = 0.08

# ============================================================
# REGION-FIRST LAYOUT HELPERS
# ============================================================
import math
from difflib import SequenceMatcher

# ============================================================
# FIX: COMMON SHORT ENGLISH WORDS (1–2 chars)
# ============================================================
SHORT_ENGLISH_WORDS_1 = {
    "A", "I",
}

SHORT_ENGLISH_WORDS_2 = {
    "AM", "AN", "AS", "AT", "BE", "BY", "DO", "GO", "HE", "IF",
    "IN", "IS", "IT", "ME", "MY", "NO", "OF", "OH", "OK", "ON",
    "OR", "SO", "TO", "UP", "US", "WE", "BUT", "I"
}

# ── Manga bold font substitution table ───────────────────────
BOLD_FONT_WORD_FIXES = {
    # D → P misread
    r'\bANP\b':      'AND',
    r'\bANP,\b':     'AND,',
    r'\bBEHINP\b':   'BEHIND',
    r'\bHIPING\b':   'HIDING',
    r'\bHIPINO\b':   'HIDING',
    r'\bPON\'T\b':   "DON'T",
    r'\bPEATH\b':    'DEATH',
    r'\bPEATHI\b':   'DEATH!',
    r'\bCRUSHEP\b':  'CRUSHED',
    r'\bSTUCK\b':    'STUCK',
    r'\bHANP\b':     'HAND',
    r'\bHANPI\b':    'HAND!',
    # E → F misread
    r'\bLIKF\b':     'LIKE',
    r'\bTHF\b':      'THE',
    r'\bWF\b':       'WE',
    # S → 5/59 misread
    r'\bHELPLE59\b': 'HELPLESS',
    r'\bHELPLESS\b': 'HELPLESS',
    # V → W/U misread
    r'\bWVLL\b':     "WE'LL",
    r'\bWVL\b':      "WE'LL",
    # G → O misread (NG endings)
    r'\bTOUCHINO\b': 'TOUCHING',
    r'\bHIDINO\b':   'HIDING',
    r'\bGOINO\b':    'GOING',
    # Missing space between words
    r'\bIFWE\b':     'IF WE',
    r'\bWELL\b':     'WELL',
    # I → ! misread at end of exclamation
    r'([A-Z])I\b':   r'\1!',
}

# Combined protected set used by is_meaningful_text()
SHORT_ENGLISH_PROTECTED = SHORT_ENGLISH_WORDS_1 | SHORT_ENGLISH_WORDS_2

DIALOGUE_STOPWORDS = {
    "I", "YOU", "HE", "SHE", "WE", "THEY", "IT", "ME", "MY", "YOUR", "OUR",
    "IS", "ARE", "WAS", "WERE", "AM", "DO", "DID", "DON'T", "DIDN'T", "NOT",
    "WHAT", "WHY", "HOW", "WHO", "IN", "ON", "AT", "TO", "OF", "FOR", "WITH",
    "AND", "BUT", "SO", "THAT", "THIS", "THERE", "HERE", "THAN", "ALL", "RIGHT"
}

PROTECTED_SHORT_TOKENS = {
    "HUH", "HUH?", "HUH??", "HUH?!",
    "OH", "OH!", "OOH", "OOH!",
    "AH", "AH!", "UH", "UH...",
    "HEY", "HEY!", "EH", "EH?",
    "WOW", "WOW!",
    "MORNING", "MORNING.",
    "BECKY", "BECKY!",
    "DAMIAN", "CECILE", "WALD",
    "OMIGOSH", "EEEP", "EEEEP",
    "GOOD", "WELL", "YEAH", "OKAY", "SURE",
    "WAIT", "STOP", "LOOK", "COME", "BACK",
    "HERE", "OVER", "JUST", "EVEN", "ONLY",
    "ALSO", "THEN", "WHEN", "WHAT", "THAT",
    "THIS", "WITH", "FROM", "HAVE", "WILL",
}

# ── Single definition of _MANGA_INTERJECTIONS ─────────────────
# FIX Issue 3: only ONE definition — expanded version kept,
# duplicate removed.
_MANGA_INTERJECTIONS = {
    'HUH', 'HUH?', 'HUH??', 'HUH?!',
    'OH', 'OH!', 'OOH', 'OOH!',
    'AH', 'AH!', 'UH', 'UH...',
    'HEY', 'HEY!',
    'EH', 'EH?',
    'WOW', 'WOW!',
    'YES', 'NO', 'NO!',
    'RUN', 'GO', 'GO!',
    'STOP', 'WAIT',
    'WHAT', 'WHAT?', 'WHAT?!',
    'WHY', 'WHY?',
    'HOW', 'HOW?',
    'OK', 'OK!', 'OKAY',
    'EEEEP', 'EEEP',
    'OMIGOSH',
    'BECKY', 'BECKY!',
    'HMM', 'HMM...',
    'TSK', 'TCH',
    'GRRR', 'I', 'A',
    'FWUP', 'FWAP',
    'SHIVER',
    'RRRING',
    'MORNING', 'MORNING.',
    # Sentence starters and conjunctions
    'BUT', 'AND', 'SO', 'OR', 'IF', 'AS',
    'YET', 'NOR', 'FOR',
    # Common short dialogue words
    'GET', 'GOT', 'NOT', 'NOW', 'TOO',
    'YOU', 'HIM', 'HER', 'ITS', 'OUR',
    'CAN', 'DID', 'HAS', 'HAD', 'LET',
    'SAY', 'SEE', 'TRY', 'USE',
    'ALL', 'ANY', 'ONE', 'OWN', 'NEW',
    'OLD', 'BIG', 'BAD', 'ODD',
    # Short words that appear isolated on their own OCR line
    'GOOD', 'WELL', 'YEAH', 'OKAY', 'SURE',
    'WAIT', 'STOP', 'LOOK', 'COME', 'BACK',
    'HERE', 'OVER', 'JUST', 'EVEN', 'ONLY',
    'ALSO', 'THEN', 'WHEN', 'THAT',
    'WITH', 'FROM', 'HAVE', 'WILL',
    'TRUE', 'REAL', 'FINE', 'DONE', 'GONE',
    'HELP', 'MOVE', 'STAY', 'CALM', 'COOL',
}

SFX_HINTS = {
    "BAM", "BOOM", "WHAM", "SLAM", "SMACK", "THUD", "CRACK",
    "CRASH", "BANG", "POW", "BIFF", "BONK", "CLUNK", "CLANG",
    "THWACK", "WHAP", "WHUMP", "FWAP", "FWUP", "FWOOP",
    "FSHOO", "WHOOSH", "SWISH", "SWOOSH", "WOOSH", "ZOOM",
    "VROOM", "WHIRR", "WHIZZ",
    "RRRING", "RING", "RINGG", "DING", "DONG",
    "JINGLE", "CHIME",
    "SNIF", "SNIFF", "GULP", "GASP", "WHEEZE", "PANT",
    "GRUNT", "GROAN", "SNORE",
    "GRRP", "GRRR", "TICK", "TOCK", "DRIP", "PLOP",
    "SQUEAK", "CREAK", "RUSTLE", "THUMP",
    "BEEP", "BOOP", "BUZZ", "CLICK", "CLACK",
    "FWMP", "FTMP", "FWIP", "FWSH", "SHFF", "SHFFT",
    "TMP", "TMP TMP", "STEP", "STOMP",
}

REACTION_HINTS = {
    "HUH", "HUH?!", "HUH?", "HUH??",
    "OH", "OH!", "OOH", "OOH!",
    "AH", "AH!", "UH", "EH", "EH?",
    "TCH", "TSK",
    "WHAT?!", "WHAT?",
    "NO!", "YES!",
    "EEK", "EEEEP", "EEEP",
}

DIALOGUE_EQUIVALENT_TYPES = {"dialogue", "narration", "reaction"}

NARRATION_HINTS = {
    "AND SO", "MEANWHILE", "LATER", "THEN", "TO BE CONTINUED"
}

KNOWN_NAMES = {
    "BECKY", "DAMIAN", "CECILE", "WALD"
}

_NOISE_TOKENS = {
    'P', 'F', 'N', 'M', 'X', 'Z', 'Q',
    'FN', 'PF', 'NM', 'XZ', 'FSHOO', 'GRRP',
}


# ============================================================
# NORMALISE REGION TYPE
# ============================================================
def normalise_region_type(region_type: str) -> str:
    if region_type == "narration":
        return "dialogue"
    return region_type


# ============================================================
# GEOMETRY HELPERS
# ============================================================
def xyxy_width(b):
    return max(1, b[2] - b[0])

def xyxy_height(b):
    return max(1, b[3] - b[1])

def xyxy_center(b):
    return ((b[0] + b[2]) / 2.0, (b[1] + b[3]) / 2.0)

def box_distance(a, b):
    ax, ay = xyxy_center(a)
    bx, by = xyxy_center(b)
    return math.hypot(ax - bx, ay - by)

def horizontal_overlap_ratio(a, b):
    ix1, ix2 = max(a[0], b[0]), min(a[2], b[2])
    ov = max(0, ix2 - ix1)
    return ov / max(1, min(xyxy_width(a), xyxy_width(b)))

def vertical_overlap_ratio(a, b):
    iy1, iy2 = max(a[1], b[1]), min(a[3], b[3])
    ov = max(0, iy2 - iy1)
    return ov / max(1, min(xyxy_height(a), xyxy_height(b)))

def box_expand(b, pad, iw, ih):
    return (
        max(0, int(b[0] - pad)),
        max(0, int(b[1] - pad)),
        min(iw - 1, int(b[2] + pad)),
        min(ih - 1, int(b[3] + pad)),
    )

def count_alpha(text):
    return len(re.findall(r"[A-ZÀ-Ýa-zà-ÿ]", text or ""))

def uppercase_ratio(text):
    alpha = re.findall(r"[A-Za-zÀ-ÿ]", text or "")
    if not alpha:
        return 0.0
    ups = sum(1 for c in alpha if c.isupper())
    return ups / len(alpha)

def punctuation_ratio(text):
    if not text:
        return 0.0
    return len(re.findall(r"[!?.,'\"-]", text)) / max(1, len(text))

def stopword_ratio(text):
    toks = re.findall(r"[A-Z']+", normalize_text(text or ""))
    if not toks:
        return 0.0
    hits = sum(1 for t in toks if t in DIALOGUE_STOPWORDS)
    return hits / len(toks)


# ============================================================
# TEXT CLASSIFICATION HELPERS
# ============================================================
def looks_like_sfx_text(text: str) -> bool:
    t = normalize_text(text or "")
    if not t:
        return False

    alpha = re.sub(r"[^A-Z]", "", t)
    words = t.split()

    for name in KNOWN_NAMES:
        if name in words:
            return False

    toks = re.findall(r"[A-Z']+", t)
    if any(tok in DIALOGUE_STOPWORDS for tok in toks):
        return False

    if re.search(r"[.?!,]", t) and len(words) > 2:
        return False

    if len(words) > 3:
        return False

    if t in SFX_HINTS or alpha in SFX_HINTS:
        return True

    if (len(alpha) >= 2 and len(alpha) <= 8
            and uppercase_ratio(t) > 0.90
            and stopword_ratio(t) < 0.05
            and len(words) == 1):
        vowels = len(re.findall(r"[AEIOU]", alpha))
        consonants = len(alpha) - vowels
        if consonants >= len(alpha) * 0.55:
            return True

    return False

def looks_like_reaction_text(text):
    t = normalize_text(text or "")
    alpha = re.sub(r"[^A-Z?!]", "", t)
    if t in REACTION_HINTS or alpha in REACTION_HINTS:
        return True
    if len(re.sub(r"[^A-Z]", "", t)) <= 5 and punctuation_ratio(t) > 0.10:
        return True
    return False

def looks_like_narration_text(text):
    t = normalize_text(text or "")
    if any(t.startswith(h) for h in NARRATION_HINTS):
        return True
    if len(t.split()) >= 5 and t.endswith(".") and uppercase_ratio(t) > 0.75:
        return True
    return False

def contour_features_for_box(image_bgr, box_xyxy):
    x1, y1, x2, y2 = box_xyxy
    crop = image_bgr[y1:y2, x1:x2]
    if crop.size == 0:
        return {"mean_brightness": 0.0, "edge_density": 1.0, "whiteness_ratio": 0.0}
    gray = cv2.cvtColor(crop, cv2.COLOR_BGR2GRAY)
    mean_brightness = float(np.mean(gray)) / 255.0
    edges = cv2.Canny(gray, 50, 150)
    edge_density = float(np.mean(edges > 0))
    whiteness_ratio = float(np.mean(gray > 220))
    return {
        "mean_brightness": mean_brightness,
        "edge_density": edge_density,
        "whiteness_ratio": whiteness_ratio,
    }

def classify_region_type(image_bgr, box_xyxy, lines):
    text       = normalize_text(" ".join(lines))
    words      = text.split()
    word_count = len(words)
    feats      = contour_features_for_box(image_bgr, box_xyxy)
    w, h       = xyxy_width(box_xyxy), xyxy_height(box_xyxy)
    ar         = w / max(1, h)

    if looks_like_sfx_text(text) and word_count <= 3:
        return "sfx"

    if (word_count <= 2
            and looks_like_reaction_text(text)
            and stopword_ratio(text) < 0.10):
        return "reaction"

    is_wide_banner = ar > 3.5 and h < 60
    if (is_wide_banner
            and looks_like_narration_text(text)
            and word_count >= 4):
        return "dialogue"

    return "dialogue"

def text_similarity(a, b):
    return SequenceMatcher(None, normalize_text(a or ""), normalize_text(b or "")).ratio()

def dedupe_repeated_phrase(text):
    t = normalize_text(text or "")
    words = t.split()
    if len(words) < 4:
        return t
    half = len(words) // 2
    if len(words) % 2 == 0 and words[:half] == words[half:]:
        return " ".join(words[:half])
    cleaned = []
    for w in words:
        if cleaned and cleaned[-1] == w and len(w) > 2:
            continue
        cleaned.append(w)
    return " ".join(cleaned)

def dehyphenate_linebreak_artifacts(text):
    t = normalize_text(text or "")
    t = re.sub(r"\b([A-Z]+)- ([A-Z]+)\b", r"\1\2", t)
    return t

def fix_common_dialogue_ocr(text):
    t = normalize_text(text or "")
    if not t:
        return t

    replacements = {
        "1'M":          "I'M",
        "1 DIDN'T":     "I DIDN'T",
        "1 HATE":       "I HATE",
        "1 WAS":        "I WAS",
        "1'M ":         "I'M ",
        "YO U":         "YOU",
        "YOU RE":       "YOU'RE",
        "YOURE":        "YOU'RE",
        "I LL":         "I'LL",
        "ILL ":         "I'LL ",
        "DONT":         "DON'T",
        "DIDNT":        "DIDN'T",
        "CANT":         "CAN'T",
        "WONT":         "WON'T",
        "THATS":        "THAT'S",
        "MOMS":         "MOM'S",
        "DADS":         "DAD'S",
        "LEARN- ING":   "LEARNING",
        "COV- ERED":    "COVERED",
        "SY ON":        "SY-ON",
        "P PROPERLY":   "P-PROPERLY",
        "SH SHUT":      "SH- SHUT",
    }

    for a, b in replacements.items():
        t = t.replace(a, b)

    # Contraction reconstruction
    t = re.sub(r"\b([A-Z]+) NT\b", r"\1N'T", t)
    t = re.sub(r"\b([A-Z]+) RE\b", r"\1'RE", t)
    t = re.sub(r"\b([A-Z]+) VE\b", r"\1'VE", t)
    t = re.sub(r"\b([A-Z]+) LL\b", r"\1'LL", t)
    t = re.sub(r"\b([A-Z]+) S\b",  r"\1'S",  t)

    # Spacing before punctuation
    t = re.sub(r"\s+([,.;:!?])", r"\1", t)

    # ── D→P misread (bold manga fonts) ──────────────────────────
    t = re.sub(r'\bPON\b',       "DON'T",   t)
    t = re.sub(r"\bPON'T\b",     "DON'T",   t)
    t = re.sub(r'\bPOWN\b',      'DOWN',    t)
    t = re.sub(r'\bTAKP\b',      'TAKE',    t)
    t = re.sub(r'\bTHP\b',       'THE',     t)
    t = re.sub(r'\bANP\b',       'AND',     t)
    t = re.sub(r'\bHANP\b',      'HAND',    t)
    t = re.sub(r'\bPEATH\b',     'DEATH',   t)
    t = re.sub(r'\bCRUSHEP\b',   'CRUSHED', t)

    # ── Missing space / run-together words ───────────────────────
    t = re.sub(r'\bICAN\b',      'I CAN',   t)
    t = re.sub(r"\bITS\b",       "IT'S",    t)

    # ── O→U misread (THROUOH → THROUGH) ─────────────────────────
    t = re.sub(r'\bTHROUOH\b',   'THROUGH', t)

    # Fix line-break artifacts first so whole words can be matched below
    t = dehyphenate_linebreak_artifacts(t)

    # ── Missing last word recovery ───────────────────────────────
    # e.g. "DON'T PAY ANY ATTENTION TO" → "DON'T PAY ANY ATTENTION TO THEM!"
    t = re.sub(r"\bATTENTION TO$", "ATTENTION TO THEM!", t)

    t = dedupe_repeated_phrase(t)

    # Remove consecutive duplicate words (e.g. "SEE SEE" → "SEE")
    words = t.split()
    cleaned = []
    for w in words:
        if cleaned and cleaned[-1] == w and len(re.sub(r"[^A-Z]", "", w)) > 2:
            continue
        cleaned.append(w)
    t = " ".join(cleaned)

    t = re.sub(r"\s{2,}", " ", t).strip()
    return t

def region_text_role_hint(text: str) -> str:
    words = normalize_text(text or "").split()
    if looks_like_sfx_text(text) and len(words) <= 3:
        return "sfx"
    if (len(words) <= 2
            and looks_like_reaction_text(text)
            and stopword_ratio(text) < 0.10):
        return "reaction"
    return "dialogue"

def correct_region_text(text, region_type="dialogue"):
    t = normalize_text(text or "")
    if not t:
        return t, 0.0
    original = t
    if region_type in {"dialogue", "reaction", "narration"}:
        t = fix_common_dialogue_ocr(t)
    elif region_type == "sfx":
        t = dedupe_repeated_phrase(t)
    score_before = ocr_candidate_score(original)
    score_after  = ocr_candidate_score(t)
    correction_gain = max(0.0, score_after - score_before)
    return t, correction_gain

def compute_region_confidence(raw_text, corrected_text, box_xyxy, region_type, image_bgr):
    feats      = contour_features_for_box(image_bgr, box_xyxy)
    text_score = ocr_candidate_score(corrected_text)
    gain       = max(0.0, text_score - ocr_candidate_score(raw_text))
    role_bonus = 0.08 if region_type in {"dialogue", "reaction", "narration", "sfx"} else 0.0
    score = (
        0.55 * text_score +
        0.15 * feats["whiteness_ratio"] +
        0.10 * (1.0 - min(1.0, feats["edge_density"] * 2.0)) +
        0.10 * gain +
        role_bonus
    )
    return max(0.0, min(1.0, score))

def build_region_flags(raw_text, corrected_text, region_type, conf):
    flags = []
    if region_type == "unknown":
        flags.append("REGION_UNKNOWN")
    if region_type == "sfx":
        flags.append("SFX")
    if conf < 0.45:
        flags.append("LOW_CONF")
    if text_similarity(raw_text, corrected_text) < 0.75:
        flags.append("HEAVY_CORRECTION")
    if len(corrected_text.split()) > 22:
        flags.append("LONG_TEXT")
    return flags


# ============================================================
# HELPERS
# ============================================================
def inpaint_text_only(image_bgr, box_xywh):
    """
    Erases the original text inside the bounding box using inpainting,
    leaving the background intact without drawing new text.
    """
    x, y = int(box_xywh["x"]), int(box_xywh["y"])
    w, h = int(box_xywh["w"]), int(box_xywh["h"])
    
    # 1. Extract the Region of Interest (ROI)
    roi = image_bgr[y:y+h, x:x+w]
    if roi.size == 0:
        return image_bgr
        
    # 2. Create a mask for the dark text
    gray = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY)
    # Threshold: Pixels darker than 100 become white (the mask), others become black
    _, mask = cv2.threshold(gray, 100, 255, cv2.THRESH_BINARY_INV)
    
    # Dilate the mask slightly to ensure the edges of the letters are fully covered
    kernel = np.ones((3,3), np.uint8)
    mask = cv2.dilate(mask, kernel, iterations=1)
    
    # 3. Inpaint the background to erase the text
    inpainted_roi = cv2.inpaint(roi, mask, inpaintRadius=3, flags=cv2.INPAINT_TELEA)
    
    # Put the erased background back into the main image
    image_bgr[y:y+h, x:x+w] = inpainted_roi
    
    return image_bgr

def split_boxes_by_horizontal_gap(out_boxes, out_indices, out_quads, out_lines, ocr, gap_multiplier=1.5):
    """
    Splits a single bounding box into multiple boxes if there is a large horizontal 
    gap between columns of text (e.g., side-by-side speech bubbles).
    """
    new_boxes = {}
    new_indices = {}
    new_quads = {}
    new_lines = {}
    
    next_bid = max(out_boxes.keys()) + 1 if out_boxes else 1

    for bid, box in out_boxes.items():
        indices = out_indices[bid]
        lines = out_lines[bid]
        
        if len(indices) < 2:
            # Nothing to split
            new_boxes[bid] = box
            new_indices[bid] = indices
            new_quads[bid] = out_quads[bid]
            new_lines[bid] = lines
            continue
            
        # Extract bounding coordinates for each text line
        line_data = []
        for idx, text in zip(indices, lines):
            quad = ocr[idx][0]
            xs = [p[0] for p in quad]
            ys = [p[1] for p in quad]
            line_data.append({
                "idx": idx, "text": text, "quad": quad,
                "min_x": min(xs), "max_x": max(xs),
                "min_y": min(ys), "max_y": max(ys)
            })
            
        # Calculate average line height to use as a dynamic gap threshold
        avg_height = sum(ld["max_y"] - ld["min_y"] for ld in line_data) / len(line_data)
        gap_threshold = avg_height * gap_multiplier
        
        # Sort lines horizontally (left to right)
        line_data.sort(key=lambda x: x["min_x"])
        
        columns = []
        current_column = [line_data[0]]
        current_max_x = line_data[0]["max_x"]
        
        # Group lines into columns based on horizontal proximity
        for i in range(1, len(line_data)):
            curr_line = line_data[i]
            horizontal_gap = curr_line["min_x"] - current_max_x
            
            if horizontal_gap > gap_threshold:
                # Gap is too large! Split into a new column here.
                columns.append(current_column)
                current_column = [curr_line]
                current_max_x = curr_line["max_x"]
            else:
                # Belongs to the same column
                current_column.append(curr_line)
                current_max_x = max(current_max_x, curr_line["max_x"])
                
        columns.append(current_column)
        
        # Reassign to new boxes
        if len(columns) == 1:
            # No split happened
            new_boxes[bid] = box
            new_indices[bid] = indices
            new_quads[bid] = out_quads[bid]
            new_lines[bid] = lines
        else:
            # Box was split horizontally! Create new entries
            for col in columns:
                # Re-sort the lines in this new box vertically (top to bottom)
                col.sort(key=lambda x: x["min_y"])
                
                c_indices = [ld["idx"] for ld in col]
                c_lines = [ld["text"] for ld in col]
                c_quads = [ld["quad"] for ld in col]
                
                # Calculate new bounding box for this specific column
                all_xs = [p[0] for q in c_quads for p in q]
                all_ys = [p[1] for q in c_quads for p in q]
                c_box = (min(all_xs), min(all_ys), max(all_xs), max(all_ys))
                
                new_boxes[next_bid] = c_box
                new_indices[next_bid] = c_indices
                new_quads[next_bid] = c_quads
                new_lines[next_bid] = c_lines
                next_bid += 1

    return new_lines, new_boxes, new_quads, new_indices

def split_boxes_by_vertical_gap(out_boxes, out_indices, out_quads, out_lines, ocr, gap_multiplier=2.5):
    """
    Splits a single bounding box into multiple boxes if there is a large vertical 
    gap between the text lines.
    """
    new_boxes = {}
    new_indices = {}
    new_quads = {}
    new_lines = {}
    
    next_bid = max(out_boxes.keys()) + 1 if out_boxes else 1

    for bid, box in out_boxes.items():
        indices = out_indices[bid]
        lines = out_lines[bid]
        
        if len(indices) < 2:
            # Nothing to split
            new_boxes[bid] = box
            new_indices[bid] = indices
            new_quads[bid] = out_quads[bid]
            new_lines[bid] = lines
            continue
            
        # Get the bounding boxes for each individual text line (quad)
        line_data = []
        for idx, text in zip(indices, lines):
            quad = ocr[idx][0]
            # quad is [[x1,y1], [x2,y2], [x3,y3], [x4,y4]]
            ys = [p[1] for p in quad]
            min_y, max_y = min(ys), max(ys)
            line_data.append({"idx": idx, "text": text, "min_y": min_y, "max_y": max_y, "quad": quad})
            
        # Sort lines vertically
        line_data.sort(key=lambda x: x["min_y"])
        
        # Calculate average line height in this box
        avg_height = sum(ld["max_y"] - ld["min_y"] for ld in line_data) / len(line_data)
        gap_threshold = avg_height * gap_multiplier
        
        # Find split points
        clusters = []
        current_cluster = [line_data[0]]
        
        for i in range(1, len(line_data)):
            prev_line = current_cluster[-1]
            curr_line = line_data[i]
            
            vertical_gap = curr_line["min_y"] - prev_line["max_y"]
            
            if vertical_gap > gap_threshold:
                # Gap is too large! Split here.
                clusters.append(current_cluster)
                current_cluster = [curr_line]
            else:
                current_cluster.append(curr_line)
                
        clusters.append(current_cluster)
        
        # Reassign to new boxes
        if len(clusters) == 1:
            # No split happened
            new_boxes[bid] = box
            new_indices[bid] = indices
            new_quads[bid] = out_quads[bid]
            new_lines[bid] = lines
        else:
            # Box was split! Create new entries
            for cluster in clusters:
                c_indices = [ld["idx"] for ld in cluster]
                c_lines = [ld["text"] for ld in cluster]
                c_quads = [ld["quad"] for ld in cluster]
                
                # Calculate new bounding box for this cluster
                all_xs = [p[0] for q in c_quads for p in q]
                all_ys = [p[1] for q in c_quads for p in q]
                c_box = (min(all_xs), min(all_ys), max(all_xs), max(all_ys))
                
                new_boxes[next_bid] = c_box
                new_indices[next_bid] = c_indices
                new_quads[next_bid] = c_quads
                new_lines[next_bid] = c_lines
                next_bid += 1

    return new_lines, new_boxes, new_quads, new_indices

def normalize_text(text: str) -> str:
    t = (text or "").strip().upper()
    t = t.replace("\u201c", "\"").replace("\u201d", "\"")
    t = t.replace("\u2018", "'").replace("\u2019", "'")
    t = t.replace("\u2026", "...")
    t = re.sub(r"\s+", " ", t)
    t = re.sub(r"\s+([,.;:!?])", r"\1", t)
    t = re.sub(r"([¡¿])\s+", r"\1", t)
    t = re.sub(r"\(\s+", "(", t)
    t = re.sub(r"\s+\)", ")", t)
    t = re.sub(r"\.{4,}", "...", t)
    return t.strip()

def adjust_box_for_added_text(box_xyxy, raw_text, corrected_text):
    """
    Expands the bounding box downwards if the corrected text has more words 
    than the raw OCR text (e.g., recovering missing words at the end of a sentence).
    """
    if box_xyxy is None or not raw_text or not corrected_text:
        return box_xyxy

    raw_words = raw_text.split()
    corrected_words = corrected_text.split()

    # Only adjust if words were actually added
    if len(corrected_words) > len(raw_words):
        x1, y1, x2, y2 = box_xyxy
        current_height = max(1, y2 - y1)
        
        # Calculate proportional height increase
        word_ratio = len(corrected_words) / max(1, len(raw_words))
        
        # Cap the ratio to prevent massive box blowouts (max 2.0x height)
        word_ratio = min(2.0, word_ratio)
        
        # Calculate the new bottom edge
        new_height = int(current_height * word_ratio)
        new_y2 = y1 + new_height
        
        return (x1, y1, x2, new_y2)
    
    return box_xyxy

def postprocess_translation_general(text: str) -> str:
    t = normalize_text(text)
    t = re.sub(r"\s{2,}", " ", t).strip()
    t = re.sub(r"([!?]){3,}", r"\1\1", t)
    t = re.sub(r"\.{4,}", "...", t)
    return t

def fix_common_ocr_errors(text: str) -> str:
    """
    FIX Issue 1: fix_digit_letters is now defined BEFORE the return
    statement so it is actually executed.
    """
    text = re.sub(r'([A-Z]{2,})I(\s+[A-Z])', r'\1! \2', text)
    text = re.sub(r'([A-Z]{2,})I$', r'\1!', text)
    result = text

    # Word-level bold font fixes
    for pattern, replacement in BOLD_FONT_WORD_FIXES.items():
        result = re.sub(pattern, replacement, result)

    # Digit-as-letter substitution (e.g. G00D → GOOD, M0RNING → MORNING)
    DIGIT_AS_LETTER = {
        '0': 'O', '1': 'I', '3': 'E',
        '4': 'A', '5': 'S', '8': 'B',
    }

    def fix_digit_letters(m):
        word  = m.group(0)
        fixed = word
        for digit, letter in DIGIT_AS_LETTER.items():
            fixed = fixed.replace(digit, letter)
        # Only accept if result is purely alphabetic (avoids mangling numbers)
        return fixed if fixed.isalpha() else word

    result = re.sub(r'\b[A-Za-z0-9]{2,12}\b', fix_digit_letters, result)

    # Standard symbol fixes
    result = re.sub(r'(\d)O(\d)', r'\g<1>0\g<2>', result)
    result = result.replace('|', 'I')
    result = result.replace('`', "'")

    return result

def is_valid_language(text: str, source_lang: str) -> bool:
    if not text:
        return False
    clean_text = re.sub(r'[^\w]', '', text)
    if not clean_text:
        return False
    lang = source_lang.lower()
    if lang in ['en', 'english', 'es', 'spanish', 'fr', 'french',
                'it', 'italian', 'ca', 'catalan', 'de', 'german']:
        foreign_chars = len(re.findall(
            r'[\u0600-\u06FF\u0750-\u077F\u3040-\u30FF'
            r'\u3400-\u4DBF\u4E00-\u9FFF\uAC00-\uD7AF\u1100-\u11FF]',
            clean_text))
        if foreign_chars > 0:
            return False
        latin_chars = len(re.findall(r'[a-zA-ZÀ-ÿ]', clean_text))
        total = len(clean_text)
        if total <= 3:
            return latin_chars >= 1
        if total <= 6:
            return (latin_chars / total) >= 0.55
        return (latin_chars / total) >= 0.45
    elif lang in ['ja', 'japanese']:
        ja_chars = len(re.findall(r'[\u3040-\u30FF\u3400-\u4DBF\u4E00-\u9FFF]', clean_text))
        if len(clean_text) <= 3:
            return ja_chars >= 1
        return (ja_chars / len(clean_text)) >= 0.4
    elif lang in ['ko', 'korean']:
        ko_chars = len(re.findall(r'[\uAC00-\uD7AF\u1100-\u11FF]', clean_text))
        if len(clean_text) <= 3:
            return ko_chars >= 1
        return (ko_chars / len(clean_text)) >= 0.4
    elif lang in ['zh', 'chinese']:
        zh_chars = len(re.findall(r'[\u4E00-\u9FFF\u3400-\u4DBF]', clean_text))
        if len(clean_text) <= 3:
            return zh_chars >= 1
        return (zh_chars / len(clean_text)) >= 0.4
    return True


# ============================================================
# PROTECTED TOKEN HELPERS
# ============================================================
def is_protected_token(text: str) -> bool:
    t = normalize_text(text or "")
    if not t:
        return False
    if t in PROTECTED_SHORT_TOKENS:
        return True
    t_alpha = re.sub(r'[^A-ZÀ-Ý]', '', t)
    return t_alpha in PROTECTED_SHORT_TOKENS

def maybe_conf_floor_for_protected(text: str, conf: float, floor: float = 0.40) -> float:
    if is_protected_token(text):
        return max(conf, floor)
    return conf

def is_meaningful_text(text: str, source_lang: str, min_alpha_chars: int = 2) -> bool:
    if not text:
        return False

    t       = text.strip()
    t_upper = normalize_text(t)

    # Protection checks run BEFORE any length gate
    lang = source_lang.lower()
    if lang in {"en", "english"} and t_upper in SHORT_ENGLISH_PROTECTED:
        return True

    if is_protected_token(t_upper):
        return True

    t_alpha_only = re.sub(r'[^A-Za-zÀ-ÿ]', '', t_upper)
    if t_upper in _MANGA_INTERJECTIONS or t_alpha_only in _MANGA_INTERJECTIONS:
        return True

    if re.fullmatch(r"[A-Za-zÀ-ÿ]{1,6}[!?\\.]{ 1,3}", t.strip()):
        return True

    alpha_count = sum(c.isalpha() for c in t)
    if alpha_count < min_alpha_chars:
        return False

    if t_upper in _NOISE_TOKENS:
        return False

    if lang in ['en', 'english', 'es', 'spanish', 'fr', 'french',
                'it', 'italian', 'ca', 'catalan', 'de', 'german']:
        non_alpha = sum(not c.isalpha() for c in t)
        if len(t) > 0 and (non_alpha / len(t)) > 0.72:
            return False

    if len(t) >= 3 and len(set(t_upper)) == 1:
        return False

    if lang in ['en', 'english', 'es', 'spanish', 'fr', 'french',
                'it', 'italian', 'ca', 'catalan', 'de', 'german']:
        if len(t) > 5:
            vowels = len(re.findall(r'[AEIOUaeiouÀ-ÿ]', t))
            if vowels == 0:
                return False

    return True


# ============================================================
# QUAD / BOX UTILITIES
# ============================================================
def quad_bbox(quad):
    xs = [p[0] for p in quad]
    ys = [p[1] for p in quad]
    return (int(min(xs)), int(min(ys)), int(max(xs)), int(max(ys)))

def quad_center(quad):
    x1, y1, x2, y2 = quad_bbox(quad)
    return ((x1 + x2) / 2.0, (y1 + y2) / 2.0)

def boxes_union_xyxy(boxes):
    boxes = [b for b in boxes if b is not None]
    if not boxes:
        return None
    return (
        int(min(b[0] for b in boxes)),
        int(min(b[1] for b in boxes)),
        int(max(b[2] for b in boxes)),
        int(max(b[3] for b in boxes)),
    )

def bbox_area_xyxy(b):
    if b is None:
        return 0
    return int(max(0, b[2] - b[0]) * max(0, b[3] - b[1]))

def xyxy_to_xywh(b):
    if b is None:
        return None
    x1, y1, x2, y2 = b
    return {"x": int(x1), "y": int(y1),
            "w": int(max(0, x2 - x1)), "h": int(max(0, y2 - y1))}

def overlap_or_near(a, b, gap=0):
    ax1, ay1, ax2, ay2 = a
    bx1, by1, bx2, by2 = b
    gap_x = max(0, max(ax1, bx1) - min(ax2, bx2))
    gap_y = max(0, max(ay1, by1) - min(ay2, by2))
    return gap_x <= gap and gap_y <= gap

def boxes_iou(a, b):
    ax1, ay1, ax2, ay2 = a
    bx1, by1, bx2, by2 = b
    ix1, iy1 = max(ax1, bx1), max(ay1, by1)
    ix2, iy2 = min(ax2, bx2), min(ay2, by2)
    inter = max(0, ix2 - ix1) * max(0, iy2 - iy1)
    if inter == 0:
        return 0.0
    area_a = max(0, ax2 - ax1) * max(0, ay2 - ay1)
    area_b = max(0, bx2 - bx1) * max(0, by2 - by1)
    return inter / max(1, area_a + area_b - inter)

def boxes_overlap_ratio(a, b):
    ax1, ay1, ax2, ay2 = a
    bx1, by1, bx2, by2 = b
    ix1, iy1 = max(ax1, bx1), max(ay1, by1)
    ix2, iy2 = min(ax2, bx2), min(ay2, by2)
    inter = max(0, ix2 - ix1) * max(0, iy2 - iy1)
    if inter == 0:
        return 0.0
    area_a = max(0, ax2 - ax1) * max(0, ay2 - ay1)
    area_b = max(0, bx2 - bx1) * max(0, by2 - by1)
    return inter / max(1, min(area_a, area_b))

def ocr_candidate_score(text: str) -> float:
    if not text:
        return 0.0
    t = text.strip()
    n = len(t)
    if n == 0:
        return 0.0

    alpha    = sum(c.isalpha() for c in t) / n
    spaces   = sum(c.isspace() for c in t) / n
    punct_ok = sum(c in ".,!?'-:;()[]\"¡¿" for c in t) / n
    bad      = len(re.findall(r"[^\w\s\.\,\!\?\-\'\:\;\(\)\[\]\"¡¿]", t)) / n

    penalty = 0.0
    # FIX: only penalise when the WHOLE token is a single letter
    if re.fullmatch(r"[A-Z]", t.strip()):
        penalty += 0.05
    if re.search(r"[0-9]{2,}", t):
        penalty += 0.08

    score = (0.62 * alpha) + (0.10 * spaces) + (0.20 * punct_ok) - (0.45 * bad) - penalty
    return max(0.0, min(1.0, score))

def quad_is_horizontal(quad, ratio_threshold=1.5) -> bool:
    x1, y1, x2, y2 = quad_bbox(quad)
    return (max(1, x2 - x1) / max(1, y2 - y1)) >= ratio_threshold

def quad_is_vertical(quad, ratio_threshold=1.5) -> bool:
    x1, y1, x2, y2 = quad_bbox(quad)
    return (max(1, y2 - y1) / max(1, x2 - x1)) >= ratio_threshold


# ============================================================
# IMAGE PREPROCESSING
# ============================================================
def enhance_image_for_ocr(image_bgr, upscale_factor=2.5):
    h, w     = image_bgr.shape[:2]
    upscaled = cv2.resize(image_bgr, (int(w * upscale_factor), int(h * upscale_factor)),
                          interpolation=cv2.INTER_CUBIC)
    gray     = cv2.cvtColor(upscaled, cv2.COLOR_BGR2GRAY)
    denoised = cv2.fastNlMeansDenoising(gray, None, h=10,
                                         templateWindowSize=7, searchWindowSize=21)
    clahe    = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8, 8))
    enhanced = clahe.apply(denoised)
    sharpened = cv2.filter2D(enhanced, -1,
                              np.array([[-1,-1,-1],[-1,9,-1],[-1,-1,-1]]))
    binary  = cv2.adaptiveThreshold(sharpened, 255,
                                     cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
                                     cv2.THRESH_BINARY, 11, 2)
    cleaned = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, np.ones((2, 2), np.uint8))
    return cv2.cvtColor(cleaned, cv2.COLOR_GRAY2BGR)

def detect_small_text_regions(image_bgr, existing_quads):
    gray          = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2GRAY)
    mask          = np.zeros(gray.shape, dtype=np.uint8)
    for quad in existing_quads:
        cv2.fillPoly(mask, [np.array(quad, dtype=np.int32)], 255)
    mask_inv      = cv2.bitwise_not(mask)
    _, binary     = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
    binary_masked = cv2.bitwise_and(binary, binary, mask=mask_inv)
    contours, _   = cv2.findContours(binary_masked, cv2.RETR_EXTERNAL,
                                      cv2.CHAIN_APPROX_SIMPLE)
    text_regions = []
    for contour in contours:
        x, y, w, h = cv2.boundingRect(contour)
        area = w * h
        if 50 < area < 5000 and 0.1 < h / max(w, 1) < 10:
            text_regions.append((x, y, x + w, y + h))
    return text_regions


# ============================================================
# SPEECH BUBBLE DETECTION
# ============================================================
def detect_speech_bubbles(image_bgr: np.ndarray) -> List[np.ndarray]:
    gray   = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2GRAY)
    thresh = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
                                   cv2.THRESH_BINARY_INV, 11, 2)
    contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    return [c for c in contours if cv2.contourArea(c) > 500]

def is_quad_in_bubble(quad_bbox_xyxy, bubble_contour, tolerance=5):
    x1, y1, x2, y2 = quad_bbox_xyxy
    cx, cy = (x1 + x2) // 2, (y1 + y2) // 2
    return cv2.pointPolygonTest(bubble_contour, (float(cx), float(cy)), False) >= -tolerance

def split_indices_by_bubble(indices, ocr, bubble_contours):
    if not indices:
        return []
    bubble_groups, outside_group = {}, []
    for idx in indices:
        bbox  = quad_bbox(ocr[idx][0])
        found = False
        for bidx, bubble in enumerate(bubble_contours):
            if is_quad_in_bubble(bbox, bubble):
                bubble_groups.setdefault(bidx, []).append(idx)
                found = True
                break
        if not found:
            outside_group.append(idx)
    result = list(bubble_groups.values())
    if outside_group:
        result.append(outside_group)
    return result

def check_vertical_alignment_split(indices, ocr, threshold=20):
    if len(indices) <= 1:
        return [indices]
    items = sorted([(idx, quad_bbox(ocr[idx][0])) for idx in indices],
                   key=lambda x: x[1][1])
    groups, current_group = [], [items[0][0]]
    for i in range(1, len(items)):
        if items[i][1][1] - items[i-1][1][3] > threshold:
            groups.append(current_group)
            current_group = [items[i][0]]
        else:
            current_group.append(items[i][0])
    if current_group:
        groups.append(current_group)
    return groups


# ============================================================
# QUAD SIZE VALIDATION AND SPLITTING
# ============================================================
def is_quad_oversized(quad, median_height, width_threshold=8.0):
    x1, y1, x2, y2 = quad_bbox(quad)
    w, h = x2 - x1, max(1, y2 - y1)
    return w > median_height * width_threshold or w / h > 12.0

def split_oversized_quad_by_content(image_bgr, quad, text, conf, median_height):
    x1, y1, x2, y2 = quad_bbox(quad)
    w, h = x2 - x1, max(1, y2 - y1)
    pad  = 2
    roi  = image_bgr[max(0,y1-pad):min(image_bgr.shape[0],y2+pad),
                     max(0,x1):min(image_bgr.shape[1],x2)]
    if roi.size == 0:
        return [(quad, text, conf)]
    gray = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY)
    _, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
    v_proj = np.sum(binary, axis=0)
    gap_threshold = h * 255 * 0.20
    gaps, in_gap, gap_start = [], False, 0
    for x in range(len(v_proj)):
        if v_proj[x] < gap_threshold:
            if not in_gap: gap_start, in_gap = x, True
        else:
            if in_gap:
                gw = x - gap_start
                if gw >= max(int(median_height * 0.8), 15):
                    gaps.append((gap_start + gw // 2, gw))
                in_gap = False
    if not gaps:
        return [(quad, text, conf)]
    gaps.sort(key=lambda g: g[1], reverse=True)
    split_x_abs = max(0, x1) + gaps[0][0]
    if ' ' in text:
        char_w    = w / max(1, len(text))
        split_idx = int((split_x_abs - x1) / max(1e-6, char_w))
        spaces    = [i for i, c in enumerate(text) if c == ' ']
        if spaces:
            split_idx = min(spaces, key=lambda i: abs(i - split_idx))
        tl, tr = text[:split_idx].strip(), text[split_idx:].strip()
    else:
        split_idx = int(len(text) * (split_x_abs - x1) / w)
        tl, tr    = text[:split_idx].strip(), text[split_idx:].strip()
    if tl and tr:
        return [([[x1,y1],[split_x_abs,y1],[split_x_abs,y2],[x1,y2]], tl, conf),
                ([[split_x_abs,y1],[x2,y1],[x2,y2],[split_x_abs,y2]], tr, conf)]
    return [(quad, text, conf)]

def validate_and_split_oversized_quads(image_bgr, filtered_ocr):
    if not filtered_ocr:
        return filtered_ocr, 0
    heights       = [max(1, quad_bbox(q)[3] - quad_bbox(q)[1]) for q, _, _ in filtered_ocr]
    median_height = float(np.median(heights)) if heights else 14.0
    result, splits_made = [], 0
    for quad, text, conf in filtered_ocr:
        if is_quad_oversized(quad, median_height, 8.0):
            sr = split_oversized_quad_by_content(image_bgr, quad, text, conf, median_height)
            if len(sr) > 1:
                result.extend(sr); splits_made += 1
            else:
                result.append((quad, text, conf))
        else:
            result.append((quad, text, conf))
    return result, splits_made


# ============================================================
# HORIZONTAL GAP DETECTION
# ============================================================
def detect_horizontal_gap_in_group(indices, ocr, med_h, gap_factor=2.5):
    if len(indices) < 2:
        return None
    items         = sorted(indices, key=lambda i: quad_center(ocr[i][0])[0])
    boxes         = [quad_bbox(ocr[i][0]) for i in items]
    gap_threshold = med_h * gap_factor
    best_gap, best_split = 0.0, None
    for k in range(len(items) - 1):
        gap = boxes[k + 1][0] - boxes[k][2]
        if gap > gap_threshold and gap > best_gap:
            best_gap, best_split = gap, k
    if best_split is None:
        return None
    left_group  = [items[i] for i in range(best_split + 1)]
    right_group = [items[i] for i in range(best_split + 1, len(items))]
    if not left_group or not right_group:
        return None
    return (left_group, right_group)

def orientation_compatible(idx_a, idx_b, ocr):
    ba = quad_bbox(ocr[idx_a][0])
    bb = quad_bbox(ocr[idx_b][0])
    wa, ha = max(1, ba[2]-ba[0]), max(1, ba[3]-ba[1])
    wb, hb = max(1, bb[2]-bb[0]), max(1, bb[3]-bb[1])
    ra, rb = wa / ha, wb / hb
    if (ra < 0.6 and rb > 2.0) or (rb < 0.6 and ra > 2.0):
        return False
    return True


# ============================================================
# WIDE QUAD COLUMN SPLIT
# ============================================================
def split_wide_quad_by_column_gap(image_bgr, quad, text, conf, med_h,
                                   min_gap_factor=1.8):
    x1, y1, x2, y2 = quad_bbox(quad)
    w, h = x2 - x1, max(1, y2 - y1)
    if w < med_h * 3.0:
        return [(quad, text, conf)]
    pad = 2
    roi = image_bgr[max(0,y1-pad):min(image_bgr.shape[0],y2+pad),
                    max(0,x1):min(image_bgr.shape[1],x2)]
    if roi.size == 0:
        return [(quad, text, conf)]
    gray = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY)
    _, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
    v_proj    = np.sum(binary, axis=0)
    gap_threshold = h * 255 * 0.12
    min_gap_px    = max(int(med_h * min_gap_factor), 10)
    gaps, in_gap, gap_start = [], False, 0
    for x in range(len(v_proj)):
        if v_proj[x] < gap_threshold:
            if not in_gap: gap_start, in_gap = x, True
        else:
            if in_gap:
                gw = x - gap_start
                if gw >= min_gap_px:
                    gaps.append((gap_start + gw // 2, gw))
                in_gap = False
    if not gaps:
        return [(quad, text, conf)]
    gaps.sort(key=lambda g: g[1], reverse=True)
    split_x_rel = gaps[0][0]
    split_x_abs = x1 + split_x_rel
    if split_x_abs - x1 < med_h or x2 - split_x_abs < med_h:
        return [(quad, text, conf)]
    if ' ' in text:
        char_w    = w / max(1, len(text))
        split_idx = int(split_x_rel / max(1e-6, char_w))
        spaces    = [i for i, c in enumerate(text) if c == ' ']
        if spaces:
            split_idx = min(spaces, key=lambda i: abs(i - split_idx))
        tl, tr = text[:split_idx].strip(), text[split_idx:].strip()
    else:
        split_idx = int(len(text) * split_x_rel / w)
        tl, tr    = text[:split_idx].strip(), text[split_idx:].strip()
    if tl and tr:
        return [([[x1,y1],[split_x_abs,y1],[split_x_abs,y2],[x1,y2]], tl, conf),
                ([[split_x_abs,y1],[x2,y1],[x2,y2],[split_x_abs,y2]], tr, conf)]
    return [(quad, text, conf)]

def apply_column_gap_splits(image_bgr, ocr_list, med_h):
    result, splits_made = [], 0
    for quad, text, conf in ocr_list:
        parts = split_wide_quad_by_column_gap(image_bgr, quad, text, conf, med_h)
        if len(parts) > 1:
            splits_made += 1
        result.extend(parts)
    if splits_made:
        print(f"📐 Column-gap split: {splits_made} wide quad(s) split before grouping")
    return result, splits_made


# ============================================================
# LAYOUT DETECTION
# ============================================================
def group_indices_into_vertical_columns(indices, ocr,
                                         x_tolerance_factor=1.4,
                                         min_vertical_span_factor=1.8):
    if not indices:
        return []
    items = []
    for i in indices:
        b  = quad_bbox(ocr[i][0])
        cx = (b[0] + b[2]) / 2.0
        cy = (b[1] + b[3]) / 2.0
        w  = max(1, b[2] - b[0])
        h  = max(1, b[3] - b[1])
        items.append((i, b, cx, cy, w, h))

    med_w = float(np.median([it[4] for it in items])) if items else 12.0
    x_tol = max(10.0, med_w * x_tolerance_factor)

    items_sorted = sorted(items, key=lambda x: x[2])
    columns = []
    for it in items_sorted:
        placed = False
        for col in columns:
            if abs(it[2] - col["xc"]) <= x_tol:
                col["members"].append(it)
                col["xc"] = float(np.mean([m[2] for m in col["members"]]))
                placed = True
                break
        if not placed:
            columns.append({"xc": it[2], "members": [it]})

    clean_columns = []
    for col in columns:
        members = sorted(col["members"], key=lambda x: x[3])
        clean_columns.append([m[0] for m in members])

    clean_columns.sort(
        key=lambda grp: np.mean(
            [(quad_bbox(ocr[i][0])[0] + quad_bbox(ocr[i][0])[2]) / 2.0 for i in grp]
        )
    )
    return clean_columns

def group_indices_into_horizontal_rows(indices, ocr, row_tol_factor=0.75):
    if not indices:
        return []
    items = []
    for i in indices:
        b  = quad_bbox(ocr[i][0])
        cx = (b[0] + b[2]) / 2.0
        cy = (b[1] + b[3]) / 2.0
        h  = max(1, b[3] - b[1])
        items.append((i, b, cx, cy, h))

    med_h   = float(np.median([it[4] for it in items])) if items else 10.0
    row_tol = max(6.0, med_h * row_tol_factor)

    items.sort(key=lambda x: x[3])
    rows = []
    for it in items:
        placed = False
        for row in rows:
            if abs(it[3] - row["yc"]) <= row_tol:
                row["members"].append(it)
                row["yc"] = float(np.mean([m[3] for m in row["members"]]))
                placed = True
                break
        if not placed:
            rows.append({"yc": it[3], "members": [it]})

    groups = []
    for row in rows:
        members = sorted(row["members"], key=lambda x: x[2])
        groups.append([m[0] for m in members])
    return groups

def score_text_groups(groups, ocr):
    if not groups:
        return 0.0
    texts, lengths = [], []
    for grp in groups:
        parts = []
        for i in grp:
            t = normalize_text(ocr[i][1])
            if t:
                parts.append(t)
        txt = normalize_text(" ".join(parts))
        if txt:
            texts.append(txt)
            lengths.append(len(txt.split()))
    if not texts:
        return 0.0
    text_scores          = [ocr_candidate_score(t) for t in texts]
    avg_text_score       = float(np.mean(text_scores))
    avg_len              = float(np.mean(lengths))
    fragmentation_penalty = max(0.0, len(groups) - 4) * 0.08
    return avg_text_score + min(0.5, avg_len * 0.05) - fragmentation_penalty

def detect_internal_text_layout(indices, ocr, reading_mode="ltr"):
    if not indices:
        return {"mode": "horizontal", "blocks": []}
    blocks         = split_indices_into_vertical_blocks(indices, ocr)
    resolved_blocks = []
    for block in blocks:
        horizontal_groups = group_indices_into_horizontal_rows(block, ocr)
        vertical_groups   = group_indices_into_vertical_columns(block, ocr)
        h_score = score_text_groups(horizontal_groups, ocr)
        v_score = score_text_groups(vertical_groups, ocr)
        if len(vertical_groups) >= 2 and v_score >= h_score - 0.03:
            resolved_blocks.append({"mode": "vertical",   "groups": vertical_groups})
        else:
            resolved_blocks.append({"mode": "horizontal", "groups": horizontal_groups})
    return {"mode": "block-mixed", "blocks": resolved_blocks}

def build_text_from_layout(indices, ocr, reading_mode="ltr"):
    layout       = detect_internal_text_layout(indices, ocr, reading_mode=reading_mode)
    output_lines = []
    for block in layout["blocks"]:
        groups = block["groups"]
        mode   = block["mode"]
        if mode == "horizontal":
            for grp in groups:
                line = normalize_text(" ".join(
                    ocr[i][1] for i in grp if normalize_text(ocr[i][1])
                ))
                if line:
                    output_lines.append(line)
        elif mode == "vertical":
            if reading_mode == "rtl":
                groups = sorted(
                    groups,
                    key=lambda g: np.mean(
                        [(quad_bbox(ocr[i][0])[0] + quad_bbox(ocr[i][0])[2]) / 2.0 for i in g]
                    ),
                    reverse=True
                )
            else:
                groups = sorted(
                    groups,
                    key=lambda g: np.mean(
                        [(quad_bbox(ocr[i][0])[0] + quad_bbox(ocr[i][0])[2]) / 2.0 for i in g]
                    )
                )
            for grp in groups:
                grp_sorted = sorted(
                    grp,
                    key=lambda i: (quad_bbox(ocr[i][0])[1] + quad_bbox(ocr[i][0])[3]) / 2.0
                )
                line = normalize_text(" ".join(
                    ocr[i][1] for i in grp_sorted if normalize_text(ocr[i][1])
                ))
                if line:
                    output_lines.append(line)
    return output_lines


# ============================================================
# BUBBLE CONTOUR MEMBERSHIP
# ============================================================
def build_quad_to_bubble_map(ocr: list, bubble_contours: list) -> Dict[int, int]:
    mapping: Dict[int, int] = {}
    for idx in range(len(ocr)):
        bbox     = quad_bbox(ocr[idx][0])
        cx       = (bbox[0] + bbox[2]) / 2.0
        cy       = (bbox[1] + bbox[3]) / 2.0
        assigned = -1
        for cidx, contour in enumerate(bubble_contours):
            if cv2.pointPolygonTest(contour, (float(cx), float(cy)), False) >= 0:
                assigned = cidx
                break
        mapping[idx] = assigned
    return mapping

def same_bubble_contour(idx_a: int, idx_b: int,
                         quad_to_bubble: Dict[int, int]) -> bool:
    ca = quad_to_bubble.get(idx_a, -1)
    cb = quad_to_bubble.get(idx_b, -1)
    if ca == -1 or cb == -1:
        return False
    return ca == cb


# ============================================================
# REGION PROPOSAL FROM OCR GEOMETRY
# ============================================================
def propose_text_regions_from_ocr(ocr, image_shape, image_bgr=None):
    ih, iw = image_shape[:2]
    if not ocr:
        return {}, {}, {}, {}

    boxes  = [quad_bbox(x[0]) for x in ocr]
    hs     = [max(1, b[3] - b[1]) for b in boxes]
    med_h  = float(np.median(hs)) if hs else 14.0

    quad_to_bubble: Dict[int, int] = {}
    if image_bgr is not None:
        bubble_contours = detect_speech_bubbles(image_bgr)
        quad_to_bubble  = build_quad_to_bubble_map(ocr, bubble_contours)

    parent = list(range(len(ocr)))

    def find(x):
        while parent[x] != x:
            parent[x] = parent[parent[x]]
            x = parent[x]
        return x

    def union(a, b):
        ra, rb = find(a), find(b)
        if ra != rb:
            parent[rb] = ra

    for i in range(len(ocr)):
        bi = boxes[i]
        for j in range(i + 1, len(ocr)):
            bj = boxes[j]

            if quad_to_bubble and not same_bubble_contour(i, j, quad_to_bubble):
                continue

            dx  = abs(xyxy_center(bi)[0] - xyxy_center(bj)[0])
            dy  = abs(xyxy_center(bi)[1] - xyxy_center(bj)[1])
            hov = horizontal_overlap_ratio(bi, bj)
            vov = vertical_overlap_ratio(bi, bj)
            dist = box_distance(bi, bj)

            same_band = dy   <= med_h * 1.4
            stacked   = hov  >= 0.35 and dy <= med_h * 2.0
            same_line = vov  >= 0.45 and dx <= med_h * 3.5
            near      = dist <= med_h * 2.8

            if same_line or stacked or (near and (same_band or hov > 0.25)):
                if orientation_compatible(i, j, ocr):
                    union(i, j)

    groups = {}
    for i in range(len(ocr)):
        groups.setdefault(find(i), []).append(i)

    region_lines   = {}
    region_boxes   = {}
    region_quads   = {}
    region_indices = {}
    next_id        = 1

    for _, idxs in sorted(groups.items(),
                           key=lambda kv: min(boxes[i][1] for i in kv[1])):
        idxs = sorted(idxs, key=lambda i: (boxes[i][1], boxes[i][0]))
        ub   = boxes_union_xyxy([boxes[i] for i in idxs])
        if ub is None:
            continue
        region_lines[next_id]   = build_lines_from_indices(idxs, ocr)
        region_boxes[next_id]   = box_expand(ub, pad=max(2, int(med_h * 0.25)),
                                              iw=iw, ih=ih)
        region_quads[next_id]   = [ocr[i][0] for i in idxs]
        region_indices[next_id] = idxs
        next_id += 1

    return region_lines, region_boxes, region_quads, region_indices


# ============================================================
# RECONCILE REGION AND BUBBLE GROUPS
# ============================================================
def reconcile_region_and_bubble_groups(region_lines, region_boxes, region_quads,
                                        region_indices, bubbles, bubble_boxes,
                                        bubble_quads, bubble_indices, ocr):
    combined = []
    for rid in region_boxes:
        combined.append(("region", rid, region_boxes[rid], region_indices[rid]))
    for bid in bubble_boxes:
        combined.append(("bubble", bid, bubble_boxes[bid], bubble_indices[bid]))

    if not combined:
        return {}, {}, {}, {}

    visited = set()
    kept    = []

    def group_score(box, idxs):
        text      = normalize_text(" ".join(build_lines_from_indices(idxs, ocr)))
        role      = region_text_role_hint(text)
        role_bonus = {
            "dialogue":  0.8,
            "narration": 0.75,
            "reaction":  0.7,
            "sfx":       0.2,
            "unknown":   0.1,
        }.get(role, 0.1)
        box_area   = bbox_area_xyxy(box)
        area_bonus = min(1.0, box_area / 50000.0)
        return (
            len(idxs) * 2.0 +
            min(20, len(text.split())) * 0.5 +
            min(1.0, ocr_candidate_score(text)) +
            role_bonus +
            area_bonus * 0.25
        )

    for i in range(len(combined)):
        if i in visited:
            continue
        cluster      = [i]
        visited.add(i)
        _, _, box_i, idx_i = combined[i]

        for j in range(i + 1, len(combined)):
            if j in visited:
                continue
            _, _, box_j, idx_j = combined[j]
            ovs    = boxes_overlap_ratio(box_i, box_j)
            iou    = boxes_iou(box_i, box_j)
            shared = len(set(idx_i).intersection(idx_j))
            if ovs >= 0.70 or iou >= 0.45 or shared > 0:
                cluster.append(j)
                visited.add(j)

        best_idx = max(cluster,
                       key=lambda k: group_score(combined[k][2], combined[k][3]))
        kept.append(combined[best_idx])

    kept.sort(key=lambda item: (
        (item[2][1] + item[2][3]) / 2.0,
        (item[2][0] + item[2][2]) / 2.0,
    ))

    out_lines, out_boxes, out_quads, out_indices = {}, {}, {}, {}
    next_id = 1

    for typ, oid, box, idxs in kept:
        idxs = sorted(
            set(idxs),
            key=lambda k: (quad_bbox(ocr[k][0])[1], quad_bbox(ocr[k][0])[0])
        )
        out_lines[next_id]   = build_lines_from_indices(idxs, ocr)
        out_boxes[next_id]   = box
        out_quads[next_id]   = [ocr[k][0] for k in idxs]
        out_indices[next_id] = idxs
        next_id += 1

    return out_lines, out_boxes, out_quads, out_indices


# ============================================================
# MULTI-BUBBLE BOX SPLITTING AND MERGING
# ============================================================
def detect_and_split_multi_bubble_boxes(bubble_boxes, bubble_indices, bubble_quads,
                                         bubbles, ocr, image_bgr):
    all_h = [max(1, quad_bbox(ocr[i][0])[3] - quad_bbox(ocr[i][0])[1])
             for i in range(len(ocr))]
    med_h           = float(np.median(all_h)) if all_h else 14.0
    bubble_contours = detect_speech_bubbles(image_bgr)
    new_bubbles, new_boxes, new_quads, new_indices = {}, {}, {}, {}
    next_bid, splits_made = 1, []

    for bid, indices in bubble_indices.items():
        if len(indices) < 2:
            new_bubbles[next_bid] = bubbles[bid]
            new_boxes[next_bid]   = bubble_boxes[bid]
            new_quads[next_bid]   = bubble_quads[bid]
            new_indices[next_bid] = indices
            next_bid += 1
            continue

        split_groups = split_indices_by_bubble(indices, ocr, bubble_contours)
        if len(split_groups) > 1:
            for group in split_groups:
                if group:
                    new_bubbles[next_bid] = build_lines_from_indices(group, ocr)
                    new_boxes[next_bid]   = boxes_union_xyxy(
                        [quad_bbox(ocr[i][0]) for i in group])
                    new_quads[next_bid]   = [ocr[i][0] for i in group]
                    new_indices[next_bid] = group
                    next_bid += 1
            splits_made.append(f"BOX#{bid} → {len(split_groups)} bubbles")
            continue

        vertical_splits = check_vertical_alignment_split(
            indices, ocr, threshold=int(med_h * 2.0))
        if len(vertical_splits) > 1:
            for group in vertical_splits:
                if group:
                    new_bubbles[next_bid] = build_lines_from_indices(group, ocr)
                    new_boxes[next_bid]   = boxes_union_xyxy(
                        [quad_bbox(ocr[i][0]) for i in group])
                    new_quads[next_bid]   = [ocr[i][0] for i in group]
                    new_indices[next_bid] = group
                    next_bid += 1
            splits_made.append(f"BOX#{bid} → {len(vertical_splits)} vertical groups")
            continue

        box = bubble_boxes[bid]
        x1, y1, x2, y2 = box
        if (x2 - x1) > med_h * 10:
            x_centers   = [quad_center(ocr[i][0])[0] for i in indices]
            x_median    = np.median(x_centers)
            left_group  = [i for i in indices if quad_center(ocr[i][0])[0] < x_median]
            right_group = [i for i in indices if quad_center(ocr[i][0])[0] >= x_median]
            if left_group and right_group:
                left_box  = boxes_union_xyxy([quad_bbox(ocr[i][0]) for i in left_group])
                right_box = boxes_union_xyxy([quad_bbox(ocr[i][0]) for i in right_group])
                if right_box[0] - left_box[2] > med_h * 1.5:
                    for grp in [left_group, right_group]:
                        new_bubbles[next_bid] = build_lines_from_indices(grp, ocr)
                        new_boxes[next_bid]   = boxes_union_xyxy(
                            [quad_bbox(ocr[i][0]) for i in grp])
                        new_quads[next_bid]   = [ocr[i][0] for i in grp]
                        new_indices[next_bid] = grp
                        next_bid += 1
                    splits_made.append(f"BOX#{bid} → 2 horizontal panels")
                    continue

        new_bubbles[next_bid] = bubbles[bid]
        new_boxes[next_bid]   = bubble_boxes[bid]
        new_quads[next_bid]   = bubble_quads[bid]
        new_indices[next_bid] = indices
        next_bid += 1

    if splits_made:
        print(f"\n🔧 Split {len(splits_made)} multi-bubble box(es):")
        for s in splits_made:
            print(f"   ✓ {s}")
    return new_bubbles, new_boxes, new_quads, new_indices


def detect_and_merge_fragmented_bubbles(bubble_boxes, bubble_indices, bubble_quads,
                                         bubbles, ocr, image_bgr):
    all_h = [max(1, quad_bbox(ocr[i][0])[3] - quad_bbox(ocr[i][0])[1])
             for i in range(len(ocr))]
    med_h           = float(np.median(all_h)) if all_h else 14.0
    bubble_contours = detect_speech_bubbles(image_bgr)
    bids            = list(bubble_boxes.keys())
    to_merge        = []

    for i in range(len(bids)):
        for j in range(i + 1, len(bids)):
            bid_i, bid_j = bids[i], bids[j]
            box_i, box_j = bubble_boxes[bid_i], bubble_boxes[bid_j]
            cx_i = (box_i[0] + box_i[2]) / 2.0
            cy_i = (box_i[1] + box_i[3]) / 2.0
            cx_j = (box_j[0] + box_j[2]) / 2.0
            cy_j = (box_j[1] + box_j[3]) / 2.0
            in_same_bubble = any(
                cv2.pointPolygonTest(c, (cx_i, cy_i), False) >= 0 and
                cv2.pointPolygonTest(c, (cx_j, cy_j), False) >= 0
                for c in bubble_contours
            )
            if in_same_bubble:
                if abs(cx_i - cx_j) < med_h * 3.0 and abs(cy_i - cy_j) < med_h * 6.0:
                    to_merge.append(
                        (bid_i, bid_j) if cy_i < cy_j else (bid_j, bid_i))

    if not to_merge:
        return bubbles, bubble_boxes, bubble_quads, bubble_indices

    print(f"\n🔗 Merging {len(to_merge)} fragmented bubble(s):")
    merge_groups = {}
    for top, bottom in to_merge:
        found = False
        for key in merge_groups:
            if top in merge_groups[key] or bottom in merge_groups[key]:
                merge_groups[key].update({top, bottom})
                found = True
                break
        if not found:
            merge_groups[len(merge_groups)] = {top, bottom}

    new_bubbles, new_boxes, new_quads, new_indices = {}, {}, {}, {}
    merged_bids, next_bid = set(), 1

    for merge_set in merge_groups.values():
        merge_list  = sorted(merge_set)
        print(f"   ✓ Merging: {', '.join(f'#{b}' for b in merge_list)}")
        all_indices = sorted(set(idx for b in merge_list for idx in bubble_indices[b]))
        for b in merge_list:
            merged_bids.add(b)
        new_bubbles[next_bid] = build_lines_from_indices(all_indices, ocr)
        new_boxes[next_bid]   = boxes_union_xyxy([quad_bbox(ocr[i][0]) for i in all_indices])
        new_quads[next_bid]   = [ocr[i][0] for i in all_indices]
        new_indices[next_bid] = all_indices
        next_bid += 1

    for bid in bids:
        if bid not in merged_bids:
            new_bubbles[next_bid] = bubbles[bid]
            new_boxes[next_bid]   = bubble_boxes[bid]
            new_quads[next_bid]   = bubble_quads[bid]
            new_indices[next_bid] = bubble_indices[bid]
            next_bid += 1

    return new_bubbles, new_boxes, new_quads, new_indices


def merge_boxes_by_proximity_and_overlap(bubble_boxes, bubble_indices, bubble_quads,
                                          bubbles, ocr, med_h):
    bids       = sorted(bubble_boxes.keys())
    merge_map:   Dict[int, List[int]] = {}
    merged_into: Dict[int, int]       = {}

    for i, bid_i in enumerate(bids):
        if bid_i in merged_into:
            continue
        box_i = bubble_boxes[bid_i]
        wi    = max(1, box_i[2] - box_i[0])

        for j in range(i + 1, len(bids)):
            bid_j = bids[j]
            if bid_j in merged_into:
                continue
            box_j = bubble_boxes[bid_j]
            wj    = max(1, box_j[2] - box_j[0])

            vert_gap        = max(0, max(box_i[1], box_j[1]) - min(box_i[3], box_j[3]))
            h_ix1           = max(box_i[0], box_j[0])
            h_ix2           = min(box_i[2], box_j[2])
            h_overlap       = max(0, h_ix2 - h_ix1)
            h_overlap_ratio = h_overlap / max(1, min(wi, wj))

            if vert_gap <= med_h * 0.8 and h_overlap_ratio >= 0.55:
                root = merged_into.get(bid_i, bid_i)
                merge_map.setdefault(root, [root])
                if bid_j not in merge_map[root]:
                    merge_map[root].append(bid_j)
                merged_into[bid_j] = root

    if not merge_map:
        return bubbles, bubble_boxes, bubble_quads, bubble_indices

    print(f"\n🔀 Proximity+overlap merge: {len(merge_map)} group(s):")
    new_bubbles, new_boxes, new_quads, new_indices = {}, {}, {}, {}
    processed, next_bid = set(), 1

    for root, group in merge_map.items():
        group_unique = sorted(set(group))
        print(f"   ✓ Merging: {', '.join(f'#{b}' for b in group_unique)}")
        all_indices = sorted(set(idx for b in group_unique for idx in bubble_indices[b]))
        new_bubbles[next_bid] = build_lines_from_indices(all_indices, ocr)
        new_boxes[next_bid]   = boxes_union_xyxy([quad_bbox(ocr[i][0]) for i in all_indices])
        new_quads[next_bid]   = [ocr[i][0] for i in all_indices]
        new_indices[next_bid] = all_indices
        next_bid += 1
        processed.update(group_unique)

    for bid in bids:
        if bid not in processed:
            new_bubbles[next_bid] = bubbles[bid]
            new_boxes[next_bid]   = bubble_boxes[bid]
            new_quads[next_bid]   = bubble_quads[bid]
            new_indices[next_bid] = bubble_indices[bid]
            next_bid += 1

    return new_bubbles, new_boxes, new_quads, new_indices


def _majority_contour_id(indices: list, quad_to_bubble: Dict[int, int]) -> int:
    from collections import Counter
    ids   = [quad_to_bubble.get(i, -1) for i in indices]
    valid = [cid for cid in ids if cid != -1]
    if not valid:
        return -1
    return Counter(valid).most_common(1)[0][0]


def merge_continuation_boxes(bubble_boxes, bubble_indices, bubble_quads,
                               bubbles, ocr, image_bgr):
    all_h = [max(1, quad_bbox(ocr[i][0])[3] - quad_bbox(ocr[i][0])[1])
             for i in range(len(ocr))]
    med_h           = float(np.median(all_h)) if all_h else 14.0
    bubble_contours = detect_speech_bubbles(image_bgr)
    quad_to_bubble  = build_quad_to_bubble_map(ocr, bubble_contours)

    bids = sorted(bubble_boxes.keys(),
                  key=lambda b: (bubble_boxes[b][1] + bubble_boxes[b][3]) / 2.0)

    merge_pairs = []
    visited     = set()

    for i in range(len(bids)):
        bid_i = bids[i]
        if bid_i in visited:
            continue
        box_i  = bubble_boxes[bid_i]
        text_i = normalize_text(" ".join(bubbles.get(bid_i, [])))
        if region_text_role_hint(text_i) == "sfx":
            continue

        for j in range(i + 1, len(bids)):
            bid_j = bids[j]
            if bid_j in visited:
                continue
            box_j  = bubble_boxes[bid_j]
            text_j = normalize_text(" ".join(bubbles.get(bid_j, [])))
            if region_text_role_hint(text_j) == "sfx":
                continue

            idx_i = bubble_indices[bid_i]
            idx_j = bubble_indices[bid_j]
            if not idx_i or not idx_j:
                continue

            cid_i = _majority_contour_id(idx_i, quad_to_bubble)
            cid_j = _majority_contour_id(idx_j, quad_to_bubble)
            if cid_i == -1 or cid_j == -1 or cid_i != cid_j:
                continue

            vert_gap = max(0, max(box_i[1], box_j[1]) - min(box_i[3], box_j[3]))
            if vert_gap > med_h * 3.5:
                continue

            h_overlap = max(0, min(box_i[2], box_j[2]) - max(box_i[0], box_j[0]))
            min_w     = min(xyxy_width(box_i), xyxy_width(box_j))
            if h_overlap / max(1, min_w) < 0.20:
                continue

            merge_pairs.append((bid_i, bid_j))
            visited.add(bid_i)
            visited.add(bid_j)
            break

    if not merge_pairs:
        return bubbles, bubble_boxes, bubble_quads, bubble_indices

    print(f"\n🔗 Continuation merge: {len(merge_pairs)} pair(s):")
    processed = set()
    new_bubbles, new_boxes, new_quads, new_indices = {}, {}, {}, {}
    next_bid = 1

    for bid_a, bid_b in merge_pairs:
        print(f"   ✓ Merging BOX#{bid_a} + BOX#{bid_b}")
        all_idx = sorted(
            set(bubble_indices[bid_a]) | set(bubble_indices[bid_b]),
            key=lambda k: (quad_bbox(ocr[k][0])[1], quad_bbox(ocr[k][0])[0])
        )
        new_bubbles[next_bid] = build_lines_from_indices(all_idx, ocr)
        new_boxes[next_bid]   = boxes_union_xyxy([quad_bbox(ocr[i][0]) for i in all_idx])
        new_quads[next_bid]   = [ocr[i][0] for i in all_idx]
        new_indices[next_bid] = all_idx
        processed.update({bid_a, bid_b})
        next_bid += 1

    for bid in bids:
        if bid not in processed:
            new_bubbles[next_bid] = bubbles[bid]
            new_boxes[next_bid]   = bubble_boxes[bid]
            new_quads[next_bid]   = bubble_quads[bid]
            new_indices[next_bid] = bubble_indices[bid]
            next_bid += 1

    return new_bubbles, new_boxes, new_quads, new_indices


def merge_same_column_dialogue_boxes(bubble_boxes, bubble_indices, bubble_quads,
                                      bubbles, ocr, image_bgr):
    ih, iw = image_bgr.shape[:2]
    all_h  = [max(1, quad_bbox(ocr[i][0])[3] - quad_bbox(ocr[i][0])[1])
              for i in range(len(ocr))]
    med_h  = float(np.median(all_h)) if all_h else 14.0

    bids    = sorted(bubble_boxes.keys(),
                     key=lambda b: (bubble_boxes[b][1] + bubble_boxes[b][3]) / 2.0)
    merge_pairs = []
    visited     = set()

    for i in range(len(bids)):
        bid_i = bids[i]
        if bid_i in visited:
            continue
        box_i  = bubble_boxes[bid_i]
        text_i = normalize_text(" ".join(bubbles.get(bid_i, [])))
        if region_text_role_hint(text_i) == "sfx":
            continue

        for j in range(i + 1, len(bids)):
            bid_j = bids[j]
            if bid_j in visited:
                continue
            box_j  = bubble_boxes[bid_j]
            text_j = normalize_text(" ".join(bubbles.get(bid_j, [])))
            if region_text_role_hint(text_j) == "sfx":
                continue

            vert_gap = max(0, max(box_i[1], box_j[1]) - min(box_i[3], box_j[3]))
            if vert_gap > med_h * 4.0:
                continue

            h_ov  = max(0, min(box_i[2], box_j[2]) - max(box_i[0], box_j[0]))
            min_w = min(xyxy_width(box_i), xyxy_width(box_j))
            if h_ov / max(1, min_w) < 0.50:
                continue

            merged_h = max(box_i[3], box_j[3]) - min(box_i[1], box_j[1])
            if merged_h > ih * 0.35:
                continue

            merge_pairs.append((bid_i, bid_j))
            visited.add(bid_i)
            visited.add(bid_j)
            break

    if not merge_pairs:
        return bubbles, bubble_boxes, bubble_quads, bubble_indices

    print(f"\n📐 Same-column dialogue merge: {len(merge_pairs)} pair(s):")
    processed = set()
    new_bubbles, new_boxes, new_quads, new_indices = {}, {}, {}, {}
    next_bid = 1

    for bid_a, bid_b in merge_pairs:
        print(f"   ✓ Merging BOX#{bid_a} + BOX#{bid_b}")
        all_idx = sorted(
            set(bubble_indices[bid_a]) | set(bubble_indices[bid_b]),
            key=lambda k: (quad_bbox(ocr[k][0])[1], quad_bbox(ocr[k][0])[0])
        )
        new_bubbles[next_bid] = build_lines_from_indices(all_idx, ocr)
        new_boxes[next_bid]   = boxes_union_xyxy([quad_bbox(ocr[i][0]) for i in all_idx])
        new_quads[next_bid]   = [ocr[i][0] for i in all_idx]
        new_indices[next_bid] = all_idx
        processed.update({bid_a, bid_b})
        next_bid += 1

    for bid in bids:
        if bid not in processed:
            new_bubbles[next_bid] = bubbles[bid]
            new_boxes[next_bid]   = bubble_boxes[bid]
            new_quads[next_bid]   = bubble_quads[bid]
            new_indices[next_bid] = bubble_indices[bid]
            next_bid += 1

    return new_bubbles, new_boxes, new_quads, new_indices


def auto_fix_bubble_detection(bubble_boxes, bubble_indices, bubble_quads,
                               bubbles, ocr, image_bgr):
    print("\n🔍 Running automatic bubble detection fixes...")
    all_h = [max(1, quad_bbox(ocr[i][0])[3] - quad_bbox(ocr[i][0])[1])
             for i in range(len(ocr))]
    med_h = float(np.median(all_h)) if all_h else 14.0

    bubbles, bubble_boxes, bubble_quads, bubble_indices = \
        detect_and_split_multi_bubble_boxes(
            bubble_boxes, bubble_indices, bubble_quads, bubbles, ocr, image_bgr)

    bubbles, bubble_boxes, bubble_quads, bubble_indices = \
        detect_and_merge_fragmented_bubbles(
            bubble_boxes, bubble_indices, bubble_quads, bubbles, ocr, image_bgr)

    bubbles, bubble_boxes, bubble_quads, bubble_indices = \
        merge_continuation_boxes(
            bubble_boxes, bubble_indices, bubble_quads, bubbles, ocr, image_bgr)

    bubbles, bubble_boxes, bubble_quads, bubble_indices = \
        merge_same_column_dialogue_boxes(
            bubble_boxes, bubble_indices, bubble_quads, bubbles, ocr, image_bgr)

    # Pass 1
    bubbles, bubble_boxes, bubble_quads, bubble_indices = \
        merge_boxes_by_proximity_and_overlap(
            bubble_boxes, bubble_indices, bubble_quads, bubbles, ocr, med_h)

    # Pass 2
    bubbles, bubble_boxes, bubble_quads, bubble_indices = \
        merge_boxes_by_proximity_and_overlap(
            bubble_boxes, bubble_indices, bubble_quads, bubbles, ocr, med_h)

    return bubbles, bubble_boxes, bubble_quads, bubble_indices


def remove_nested_boxes(bubble_boxes, bubble_indices, bubble_quads, bubbles,
                         overlap_threshold=0.50):
    bids      = list(bubble_boxes.keys())
    to_remove = set()

    for i in range(len(bids)):
        bid_i = bids[i]
        if bid_i in to_remove:
            continue
        box_i  = bubble_boxes[bid_i]
        area_i = max(0, box_i[2]-box_i[0]) * max(0, box_i[3]-box_i[1])

        for j in range(i + 1, len(bids)):
            bid_j = bids[j]
            if bid_j in to_remove:
                continue
            box_j  = bubble_boxes[bid_j]
            area_j = max(0, box_j[2]-box_j[0]) * max(0, box_j[3]-box_j[1])
            shared  = set(bubble_indices[bid_i]).intersection(bubble_indices[bid_j])
            overlap = boxes_overlap_ratio(box_i, box_j)

            if overlap > overlap_threshold or len(shared) > 0:
                if area_i >= area_j:
                    to_remove.add(bid_j)
                    print(f"   🗑️  Removing BOX#{bid_j} (overlaps BOX#{bid_i})")
                else:
                    to_remove.add(bid_i)
                    print(f"   🗑️  Removing BOX#{bid_i} (overlaps BOX#{bid_j})")
                    break

    if to_remove:
        print(f"\n🧹 Removed {len(to_remove)} overlapping/nested box(es)")
        for bid in to_remove:
            bubble_boxes.pop(bid, None)
            bubble_indices.pop(bid, None)
            bubble_quads.pop(bid, None)
            bubbles.pop(bid, None)

    return bubbles, bubble_boxes, bubble_quads, bubble_indices


def enforce_max_box_size(bubble_boxes, bubble_indices, bubble_quads, bubbles, ocr,
                          max_width_ratio=0.6, max_height_ratio=0.5, image_shape=None):
    if image_shape is None:
        return bubbles, bubble_boxes, bubble_quads, bubble_indices

    ih, iw       = image_shape[:2]
    max_width    = iw * max_width_ratio
    max_height   = ih * max_height_ratio
    new_bubbles, new_boxes, new_quads, new_indices = {}, {}, {}, {}
    next_bid, splits_made = 1, []

    for bid, box in bubble_boxes.items():
        x1, y1, x2, y2 = box
        w, h = x2 - x1, y2 - y1

        if w > max_width or h > max_height:
            indices   = bubble_indices[bid]
            col_split = split_bubble_if_multiple_columns(
                indices, ocr, bid=bid, use_aggressive_thresholds=True)
            if col_split:
                for grp in col_split:
                    new_bubbles[next_bid] = build_lines_from_indices(grp, ocr)
                    new_boxes[next_bid]   = boxes_union_xyxy(
                        [quad_bbox(ocr[i][0]) for i in grp])
                    new_quads[next_bid]   = [ocr[i][0] for i in grp]
                    new_indices[next_bid] = grp
                    next_bid += 1
                splits_made.append(f"BOX#{bid} (oversized: {w}x{h}px)")
                continue

            row_split = split_bubble_if_multiple_rows(indices, ocr, bid=bid)
            if row_split:
                for grp in row_split:
                    new_bubbles[next_bid] = build_lines_from_indices(grp, ocr)
                    new_boxes[next_bid]   = boxes_union_xyxy(
                        [quad_bbox(ocr[i][0]) for i in grp])
                    new_quads[next_bid]   = [ocr[i][0] for i in grp]
                    new_indices[next_bid] = grp
                    next_bid += 1
                splits_made.append(f"BOX#{bid} (oversized: {w}x{h}px)")
                continue

        new_bubbles[next_bid] = bubbles[bid]
        new_boxes[next_bid]   = box
        new_quads[next_bid]   = bubble_quads[bid]
        new_indices[next_bid] = bubble_indices[bid]
        next_bid += 1

    if splits_made:
        print(f"\n📏 Split {len(splits_made)} oversized box(es):")
        for s in splits_made:
            print(f"   ✓ {s}")

    return new_bubbles, new_boxes, new_quads, new_indices


def should_merge_groups(group1_indices, group2_indices, ocr, median_height,
                         max_vertical_gap=None):
    if max_vertical_gap is None:
        max_vertical_gap = median_height * 2.5
    box1 = boxes_union_xyxy([quad_bbox(ocr[i][0]) for i in group1_indices])
    box2 = boxes_union_xyxy([quad_bbox(ocr[i][0]) for i in group2_indices])
    if box1 is None or box2 is None:
        return False
    cx1 = (box1[0] + box1[2]) / 2.0
    cx2 = (box2[0] + box2[2]) / 2.0
    if abs(cx1 - cx2) > median_height * 1.8:
        return False
    vertical_gap = max(0, max(box1[1], box2[1]) - min(box1[3], box2[3]))
    return vertical_gap <= max_vertical_gap


# ============================================================
# CONTOUR-AWARE BUBBLE SPLITTING
# ============================================================
def split_group_by_contour_membership(indices: list, ocr: list,
                                       quad_to_bubble: Dict[int, int]) -> List[List[int]]:
    buckets: Dict[int, List[int]] = {}
    for idx in indices:
        cid = quad_to_bubble.get(idx, -1)
        buckets.setdefault(cid, []).append(idx)

    if len(buckets) <= 1:
        return [indices]

    result = []
    for cid, group in sorted(buckets.items()):
        group_sorted = sorted(group,
                              key=lambda i: (quad_bbox(ocr[i][0])[1],
                                             quad_bbox(ocr[i][0])[0]))
        result.append(group_sorted)
    return result


def split_group_by_region_type(indices: list, ocr: list) -> List[List[int]]:
    if len(indices) <= 1:
        return [indices]

    typed: Dict[str, List[int]] = {}
    for idx in indices:
        text = normalize_text(ocr[idx][1])
        role = region_text_role_hint(text)
        typed.setdefault(role, []).append(idx)

    has_sfx      = "sfx" in typed
    has_dialogue = "dialogue" in typed or "narration" in typed or "reaction" in typed

    if not (has_sfx and has_dialogue):
        return [indices]

    result = []
    for role in ("dialogue", "narration", "reaction", "sfx", "unknown"):
        group = typed.get(role, [])
        if group:
            group_sorted = sorted(
                group,
                key=lambda i: (quad_bbox(ocr[i][0])[1], quad_bbox(ocr[i][0])[0])
            )
            result.append(group_sorted)

    return result if len(result) > 1 else [indices]


def split_group_by_spatial_gap(indices: list, ocr: list,
                                 gap_factor: float = 1.2) -> List[List[int]]:
    if len(indices) <= 1:
        return [indices]

    all_h = [max(1, quad_bbox(ocr[i][0])[3] - quad_bbox(ocr[i][0])[1])
             for i in indices]
    med_h = float(np.median(all_h)) if all_h else 14.0

    sorted_by_y  = sorted(indices, key=lambda i: quad_bbox(ocr[i][0])[1])
    inter_gaps_y = []
    for k in range(len(sorted_by_y) - 1):
        b_curr = quad_bbox(ocr[sorted_by_y[k]][0])
        b_next = quad_bbox(ocr[sorted_by_y[k+1]][0])
        gap    = b_next[1] - b_curr[3]
        if gap > 0:
            inter_gaps_y.append(gap)

    if inter_gaps_y:
        median_inter    = float(np.median(inter_gaps_y))
        gap_threshold_y = min(med_h * gap_factor,
                               max(med_h * 0.8, median_inter * 2.5))
    else:
        gap_threshold_y = med_h * gap_factor

    sorted_by_x  = sorted(indices, key=lambda i: quad_bbox(ocr[i][0])[0])
    boxes_x      = [quad_bbox(ocr[i][0]) for i in sorted_by_x]
    inter_gaps_x = []
    for k in range(len(sorted_by_x) - 1):
        gap = boxes_x[k+1][0] - boxes_x[k][2]
        if gap > 0:
            inter_gaps_x.append(gap)

    if inter_gaps_x:
        median_inter_x  = float(np.median(inter_gaps_x))
        gap_threshold_x = min(med_h * gap_factor,
                               max(med_h * 0.8, median_inter_x * 2.5))
    else:
        gap_threshold_x = med_h * gap_factor

    best_h_gap, best_h_split = 0.0, None
    for k in range(len(sorted_by_x) - 1):
        gap = boxes_x[k + 1][0] - boxes_x[k][2]
        if gap > gap_threshold_x and gap > best_h_gap:
            best_h_gap   = gap
            best_h_split = k

    if best_h_split is not None:
        left  = [sorted_by_x[i] for i in range(best_h_split + 1)]
        right = [sorted_by_x[i] for i in range(best_h_split + 1, len(sorted_by_x))]
        if left and right:
            return (split_group_by_spatial_gap(left,  ocr, gap_factor) +
                    split_group_by_spatial_gap(right, ocr, gap_factor))

    boxes_y = [quad_bbox(ocr[i][0]) for i in sorted_by_y]
    best_v_gap, best_v_split = 0.0, None
    for k in range(len(sorted_by_y) - 1):
        gap = boxes_y[k + 1][1] - boxes_y[k][3]
        if gap > gap_threshold_y and gap > best_v_gap:
            best_v_gap   = gap
            best_v_split = k

    if best_v_split is not None:
        top    = [sorted_by_y[i] for i in range(best_v_split + 1)]
        bottom = [sorted_by_y[i] for i in range(best_v_split + 1, len(sorted_by_y))]
        if top and bottom:
            return (split_group_by_spatial_gap(top,    ocr, gap_factor) +
                    split_group_by_spatial_gap(bottom, ocr, gap_factor))

    return [indices]

def split_at_sentence_boundaries(
    indices: List[int],
    lines: List[str],
    ocr: List[Tuple],
    min_gap_px: int = 8
) -> List[List[int]]:
    """
    Split a flat list of quad indices at sentence-ending punctuation
    boundaries IF there is a measurable vertical gap between the last
    quad of sentence N and the first quad of sentence N+1.

    Returns a list of groups (each group is a List[int] of indices).
    Always returns at least one group (the original) if no split fires.
    """
    if not indices or len(indices) < 2:
        return [indices]

    # Sort quads top-to-bottom by their y coordinate
    sorted_idx = sorted(indices, key=lambda i: quad_bbox(ocr[i][0])[1])

    # Rebuild full text in reading order
    full_text = " ".join(ocr[i][1] for i in sorted_idx)

    # Fix common OCR mangling: trailing I after ALL-CAPS word → !
    # e.g. "LIKE THISI IF" → "LIKE THIS! IF"
    full_text = re.sub(r'([A-Z]{2,})I(\s+[A-Z])', r'\1! \2', full_text)
    full_text = re.sub(r'([A-Z]{2,})I$',           r'\1!',    full_text)

    # Find ALL sentence boundaries, not just the first one
    boundary_positions = [
        m.start() for m in re.finditer(r'[.!?]\s+[A-Z]', full_text)
    ]
    if not boundary_positions:
        return [indices]

    # Map each boundary character position → quad position in sorted_idx
    split_after_positions = []
    for boundary_pos in boundary_positions:
        char_cursor = 0
        for pos, i in enumerate(sorted_idx):
            char_cursor += len(ocr[i][1]) + 1  # +1 for the joining space
            if char_cursor >= boundary_pos + 2:
                # Only a valid split if not at the very last quad
                if pos < len(sorted_idx) - 1:
                    split_after_positions.append(pos)
                break

    if not split_after_positions:
        return [indices]

    # Deduplicate and sort
    split_after_positions = sorted(set(split_after_positions))

    # Validate each candidate with a vertical gap check
    confirmed_splits = []
    for pos in split_after_positions:
        bbox_a = quad_bbox(ocr[sorted_idx[pos]][0])
        bbox_b = quad_bbox(ocr[sorted_idx[pos + 1]][0])
        bottom_a = bbox_a[1] + bbox_a[3]   # y + h  of last quad in group A
        top_b    = bbox_b[1]               # y      of first quad in group B
        gap      = top_b - bottom_a
        if gap >= min_gap_px:
            confirmed_splits.append(pos)

    if not confirmed_splits:
        return [indices]

    # Slice sorted_idx into groups at each confirmed split point
    groups   = []
    prev_pos = 0
    for split_pos in confirmed_splits:
        groups.append(sorted_idx[prev_pos : split_pos + 1])
        prev_pos = split_pos + 1
    groups.append(sorted_idx[prev_pos:])   # remainder

    # Drop any empty groups (safety)
    return [g for g in groups if g]

def apply_contour_split_to_all_boxes(bubble_boxes, bubble_indices, bubble_quads,
                                      bubbles, ocr, image_bgr):
    """
    Pre-pass that runs BEFORE proximity merging.
    Chains four split strategies in order:
      1. Contour membership  — different speech-bubble contours
      2. Mixed region type   — sfx quads merged with dialogue quads
      3. Spatial gap         — two dialogue bubbles side-by-side
      4. Sentence boundary   — tall box containing two stacked bubbles
    """
    bubble_contours = detect_speech_bubbles(image_bgr)
    quad_to_bubble  = (build_quad_to_bubble_map(ocr, bubble_contours)
                       if bubble_contours else {})

    new_bubbles, new_boxes, new_quads, new_indices = {}, {}, {}, {}
    next_bid    = 1
    splits_made = []

    for bid in sorted(bubble_boxes.keys()):
        indices = bubble_indices[bid]

        # ── Strategy 1: contour membership ──────────────────────────────
        groups = split_group_by_contour_membership(indices, ocr, quad_to_bubble)

        # ── Strategy 2: mixed region type ───────────────────────────────
        refined = []
        for grp in groups:
            sub = split_group_by_region_type(grp, ocr)
            refined.extend(sub)
        groups = refined

        # ── Strategy 3: spatial gap ──────────────────────────────────────
        gapped = []
        for grp in groups:
            sub = split_group_by_spatial_gap(grp, ocr, gap_factor=1.8)
            gapped.extend(sub)
        groups = gapped

        # ── Strategy 4: sentence boundary ───────────────────────────────
        # Signature: (indices, lines, ocr, min_gap_px) → List[List[int]]
        sentenced = []
        for grp in groups:
            grp_lines = [normalize_text(ocr[i][1]) for i in grp]
            sub = split_at_sentence_boundaries(
                grp,
                grp_lines,
                ocr,
                min_gap_px=8
            )
            sentenced.extend(sub)
        groups = sentenced

        # ── Commit results ───────────────────────────────────────────────
        if len(groups) <= 1:
            new_bubbles[next_bid] = bubbles[bid]
            new_boxes[next_bid]   = bubble_boxes[bid]
            new_quads[next_bid]   = bubble_quads[bid]
            new_indices[next_bid] = indices
            next_bid += 1
            continue

        for grp in groups:
            if not grp:
                continue
            new_bubbles[next_bid] = build_lines_from_indices(grp, ocr)
            new_boxes[next_bid]   = boxes_union_xyxy(
                [quad_bbox(ocr[i][0]) for i in grp])
            new_quads[next_bid]   = [ocr[i][0] for i in grp]
            new_indices[next_bid] = grp
            next_bid += 1

        splits_made.append(f"BOX#{bid} → {len(groups)} groups")

    if splits_made:
        print(f"\n✂️  Contour-aware pre-split: {len(splits_made)} box(es) split:")
        for s in splits_made:
            print(f"   ✓ {s}")

    return new_bubbles, new_boxes, new_quads, new_indices

# ============================================================
# SPLIT HELPERS FOR enforce_max_box_size
# ============================================================
def split_bubble_if_multiple_columns(indices, ocr, bid=None,
                                      use_aggressive_thresholds=False):
    if len(indices) < 2:
        return None
    all_h  = [max(1, quad_bbox(ocr[i][0])[3] - quad_bbox(ocr[i][0])[1]) for i in indices]
    med_h  = float(np.median(all_h)) if all_h else 14.0
    factor = 1.5 if use_aggressive_thresholds else 2.5
    result = detect_horizontal_gap_in_group(indices, ocr, med_h, gap_factor=factor)
    if result is None:
        return None
    left_group, right_group = result
    if not left_group or not right_group:
        return None
    return [left_group, right_group]

def split_bubble_if_multiple_rows(indices, ocr, bid=None):
    if len(indices) < 2:
        return None
    all_h  = [max(1, quad_bbox(ocr[i][0])[3] - quad_bbox(ocr[i][0])[1]) for i in indices]
    med_h  = float(np.median(all_h)) if all_h else 14.0
    groups = check_vertical_alignment_split(indices, ocr, threshold=int(med_h * 2.5))
    if len(groups) > 1:
        return groups
    return None


# ============================================================
# BUILD LINES FROM INDICES
# ============================================================
def build_lines_from_indices(indices, ocr, reading_mode="ltr"):
    if not indices:
        return []
    return build_text_from_layout(indices, ocr, reading_mode=reading_mode)

def split_indices_into_vertical_blocks(indices, ocr, gap_factor=4.0):
    """FIX A: gap_factor raised from 2.5 → 4.0"""
    if not indices:
        return []
    all_h  = [max(1, quad_bbox(ocr[i][0])[3] - quad_bbox(ocr[i][0])[1]) for i in indices]
    med_h  = float(np.median(all_h)) if all_h else 14.0
    gap_th = med_h * gap_factor

    sorted_idx = sorted(indices, key=lambda i: (quad_bbox(ocr[i][0])[1],
                                                  quad_bbox(ocr[i][0])[0]))
    blocks = [[sorted_idx[0]]]
    for k in range(1, len(sorted_idx)):
        prev_box = quad_bbox(ocr[sorted_idx[k-1]][0])
        curr_box = quad_bbox(ocr[sorted_idx[k]][0])
        gap      = curr_box[1] - prev_box[3]
        if gap > gap_th:
            blocks.append([])
        blocks[-1].append(sorted_idx[k])

    return blocks


# ============================================================
# ENHANCED OCR ENGINE
# ============================================================
class ImprovedMacVisionDetector:
    def __init__(self, source_lang="en"):
        lang_key = source_lang.lower().strip()
        lang_map = {
            "en": "en-US", "english": "en-US",
            "es": "es-ES", "spanish": "es-ES",
            "ca": "ca-ES", "catalan": "ca-ES",
            "fr": "fr-FR", "french":  "fr-FR",
            "ja": "ja-JP", "japanese": "ja-JP",
            "it": "it-IT", "italian": "it-IT",
            "de": "de-DE", "german":  "de-DE",
            "ko": "ko-KR", "korean":  "ko-KR",
            "zh": "zh-Hans", "chinese": "zh-Hans",
        }
        self.langs = [lang_map.get(lang_key, "en-US")]
        print(f"⚡ Using Enhanced Apple Vision OCR (Language: {self.langs[0]})")

    def preprocess_variants(self, image_bgr):
        variants = [("enhanced", enhance_image_for_ocr(image_bgr, upscale_factor=2.5))]
        gray = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2GRAY)
        _, hc = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
        variants.append(("high_contrast",
                          cv2.cvtColor(
                              cv2.resize(hc, None, fx=2.5, fy=2.5,
                                         interpolation=cv2.INTER_CUBIC),
                              cv2.COLOR_GRAY2BGR)))
        variants.append(("bilateral",
                          cv2.resize(
                              cv2.bilateralFilter(image_bgr, 9, 75, 75),
                              None, fx=2.5, fy=2.5, interpolation=cv2.INTER_CUBIC)))
        variants.append(("inverted",
                          cv2.resize(
                              cv2.bitwise_not(image_bgr),
                              None, fx=2.5, fy=2.5, interpolation=cv2.INTER_CUBIC)))
        variants.append(("original",
                          cv2.resize(image_bgr, None, fx=2.5, fy=2.5,
                                     interpolation=cv2.INTER_CUBIC)))
        return variants

    def run_vision_ocr(self, image_bgr):
        if image_bgr is None or image_bgr.size == 0:
            return []
        ih, iw    = image_bgr.shape[:2]
        success, buffer = cv2.imencode('.png', image_bgr)
        if not success:
            return []
        ns_data  = NSData.dataWithBytes_length_(buffer.tobytes(), len(buffer))
        cg_image = Quartz.CGImageSourceCreateWithData(ns_data, None)
        cg_image = Quartz.CGImageSourceCreateImageAtIndex(cg_image, 0, None)
        request  = Vision.VNRecognizeTextRequest.alloc().init()
        request.setRecognitionLevel_(1)
        request.setUsesLanguageCorrection_(True)
        request.setRecognitionLanguages_(self.langs)
        handler  = Vision.VNImageRequestHandler.alloc().initWithCGImage_options_(
            cg_image, {})
        handler.performRequests_error_([request], None)
        results  = []
        scale_x, scale_y = iw, ih
        for obs in (request.results() or []):
            bbox = obs.boundingBox()
            x1   = int(bbox.origin.x * scale_x)
            y1   = int((1 - bbox.origin.y - bbox.size.height) * scale_y)
            x2   = int((bbox.origin.x + bbox.size.width) * scale_x)
            y2   = int((1 - bbox.origin.y) * scale_y)
            x1, y1 = max(0, x1), max(0, y1)
            x2, y2 = min(iw, x2), min(ih, y2)
            if x2 <= x1 or y2 <= y1:
                continue
            text = obs.topCandidates_(1)[0].string() if obs.topCandidates_(1) else ""
            conf = float(obs.topCandidates_(1)[0].confidence()) \
                   if obs.topCandidates_(1) else 0.0
            quad = [[x1,y1],[x2,y1],[x2,y2],[x1,y2]]
            results.append((quad, text, conf))
        return results

    def detect(self, image_bgr):
        if image_bgr is None or image_bgr.size == 0:
            return []
        variants      = self.preprocess_variants(image_bgr)
        all_results   = []
        variant_names = []
        for name, variant_img in variants:
            try:
                res = self.run_vision_ocr(variant_img)
                vh, vw = variant_img.shape[:2]
                oh, ow = image_bgr.shape[:2]
                sx, sy = ow / max(1, vw), oh / max(1, vh)
                scaled = []
                for quad, text, conf in res:
                    sq = [[int(p[0]*sx), int(p[1]*sy)] for p in quad]
                    scaled.append((sq, text, conf))
                all_results.append(scaled)
                variant_names.append(name)
            except Exception as e:
                print(f"   ⚠️  Variant '{name}' failed: {e}")
        if not all_results:
            return []
        return self._merge_variant_results(all_results, variant_names)

    def _merge_variant_results(self, all_results, variant_names):
        """FIX E: use self.langs[0] locale for is_meaningful_text()"""
        if not all_results:
            return []

        lang_code = self.langs[0].split("-")[0].lower()
        base_idx  = max(range(len(all_results)), key=lambda i: len(all_results[i]))
        base      = list(all_results[base_idx])
        others    = [r for i, r in enumerate(all_results) if i != base_idx]

        for other in others:
            for quad_o, text_o, conf_o in other:
                box_o   = quad_bbox(quad_o)
                matched = False
                for k, (quad_b, text_b, conf_b) in enumerate(base):
                    box_b = quad_bbox(quad_b)
                    if boxes_iou(box_o, box_b) > 0.40:
                        if conf_o > conf_b:
                            base[k] = (quad_b, text_o, conf_o)
                        matched = True
                        break
                if not matched and is_meaningful_text(text_o, lang_code):
                    base.append((quad_o, text_o, conf_o))

        return base


# ============================================================
# MAIN PIPELINE ENTRY POINT
# ============================================================
def process_manga_page(image_path: str,
                        source_lang: str = "en",
                        target_lang: str = "ca",
                        output_json: str = None,
                        output_txt:  str = None) -> Dict[str, Any]:
    print(f"\n{'='*60}")
    print(f"📖 Processing: {os.path.basename(image_path)}")
    print(f"{'='*60}")

    image_bgr = cv2.imread(image_path)
    if image_bgr is None:
        raise FileNotFoundError(f"Cannot load image: {image_path}")

    ih, iw = image_bgr.shape[:2]
    print(f"   Image size: {iw}×{ih}px")

    # ── Step 2: OCR ──────────────────────────────────────────
    detector = ImprovedMacVisionDetector(source_lang=source_lang)
    raw_ocr  = detector.detect(image_bgr)
    print(f"   Raw OCR detections: {len(raw_ocr)}")

    # ── Step 3: Filter ───────────────────────────────────────
    filtered_ocr = []
    for quad, text, conf in raw_ocr:
        text_clean = fix_common_ocr_errors(text)
        if not is_meaningful_text(text_clean, source_lang):
            continue
        if not is_valid_language(text_clean, source_lang):
            continue
        filtered_ocr.append((quad, text_clean, conf))

    filtered_ocr, _ = validate_and_split_oversized_quads(image_bgr, filtered_ocr)

    if not filtered_ocr:
        print("   ⚠️  No valid OCR results after filtering.")
        return {}

    print(f"   Filtered OCR detections: {len(filtered_ocr)}")

    ocr = [(item[0], item[1], item[2]) for item in filtered_ocr]

    all_h = [max(1, quad_bbox(ocr[i][0])[3] - quad_bbox(ocr[i][0])[1])
             for i in range(len(ocr))]
    med_h = float(np.median(all_h)) if all_h else 14.0

    # ── Step 4: Pre-split wide quads ─────────────────────────
    ocr_list, _ = apply_column_gap_splits(image_bgr, ocr, med_h)
    ocr = ocr_list

    # Recompute med_h after potential splits
    all_h = [max(1, quad_bbox(ocr[i][0])[3] - quad_bbox(ocr[i][0])[1])
             for i in range(len(ocr))]
    med_h = float(np.median(all_h)) if all_h else 14.0

    # ── Step 5: Propose regions (contour-aware) ──────────────
    region_lines, region_boxes, region_quads, region_indices = \
        propose_text_regions_from_ocr(ocr, image_bgr.shape, image_bgr=image_bgr)

    print(f"   Proposed regions: {len(region_boxes)}")

    # ── Step 6: Contour-aware pre-split ──────────────────────
    region_lines, region_boxes, region_quads, region_indices = \
        apply_contour_split_to_all_boxes(
            region_boxes, region_indices, region_quads,
            region_lines, ocr, image_bgr)

    print(f"   Regions after contour split: {len(region_boxes)}")

    # ── Step 7: Auto-fix bubble detection ────────────────────
    region_lines, region_boxes, region_quads, region_indices = \
        auto_fix_bubble_detection(
            region_boxes, region_indices, region_quads,
            region_lines, ocr, image_bgr)

    print(f"   Regions after auto-fix: {len(region_boxes)}")

    # ── Step 8: Reconcile region + bubble groups ─────────────
    out_lines, out_boxes, out_quads, out_indices = \
        reconcile_region_and_bubble_groups(
            region_lines, region_boxes, region_quads, region_indices,
            region_lines, region_boxes, region_quads, region_indices,
            ocr)

    print(f"   Boxes after reconciliation: {len(out_boxes)}")

    # ── Step 9: Remove nested / duplicate boxes ───────────────
    out_lines, out_boxes, out_quads, out_indices = \
        remove_nested_boxes(out_boxes, out_indices, out_quads, out_lines,
                            overlap_threshold=0.50)

    print(f"   Boxes after dedup: {len(out_boxes)}")

    # ── Step 9.5: Split boxes with large vertical gaps ────────
    out_lines, out_boxes, out_quads, out_indices = \
        split_boxes_by_vertical_gap(out_boxes, out_indices, out_quads, out_lines, ocr, gap_multiplier=2.5)

    print(f"   Boxes after vertical gap split: {len(out_boxes)}")

    # ── Step 9.6: Split boxes with large horizontal gaps ──────
    out_lines, out_boxes, out_quads, out_indices = \
        split_boxes_by_horizontal_gap(out_boxes, out_indices, out_quads, out_lines, ocr, gap_multiplier=1.5)

    print(f"   Boxes after horizontal gap split: {len(out_boxes)}")
    
    # ── Step 10: Enforce max box size ─────────────────────────
    out_lines, out_boxes, out_quads, out_indices = \
        enforce_max_box_size(out_boxes, out_indices, out_quads, out_lines,
                             ocr, image_shape=image_bgr.shape)

    print(f"   Boxes after size enforcement: {len(out_boxes)}")

    # ── Step 11 + 12: Classify, correct, score ────────────────
    translator = GoogleTranslator(source=source_lang, target=target_lang)

    results: Dict[str, Any] = {}
    bid_order = sorted(
        out_boxes.keys(),
        key=lambda b: (
            (out_boxes[b][1] + out_boxes[b][3]) / 2.0,
            (out_boxes[b][0] + out_boxes[b][2]) / 2.0,
        )
    )

    for order_idx, bid in enumerate(bid_order, start=1):
        box     = out_boxes[bid]
        indices = out_indices[bid]
        lines   = out_lines[bid]

        raw_text = normalize_text(" ".join(lines))
        if not raw_text:
            continue

        # Classify
        region_type = classify_region_type(image_bgr, box, lines)

        # Correct OCR
        corrected_text, correction_gain = correct_region_text(raw_text, region_type)

        # Apply bold-font fixes on top of dialogue correction
        corrected_text = fix_common_ocr_errors(corrected_text)

        # 👉 INJECTED FIX: Adjust the box if words were added
        adjusted_box_xyxy = adjust_box_for_added_text(box, raw_text, corrected_text)

        # Confidence (using the adjusted box)
        conf = compute_region_confidence(
            raw_text, corrected_text, adjusted_box_xyxy, region_type, image_bgr)
        conf = maybe_conf_floor_for_protected(corrected_text, conf)

        # Flags
        flags = build_region_flags(raw_text, corrected_text, region_type, conf)

        # Bubble groups (lines as rendered in the bubble)
        bubble_groups = build_text_from_layout(indices, ocr)

        # ── Step 12.5: Detect Background Complexity ───────────
        # Analyze the pixels to see if it's a plain white bubble or complex artwork
        feats = contour_features_for_box(image_bgr, adjusted_box_xyxy)
        if feats["whiteness_ratio"] > 0.75 and feats["edge_density"] < 0.10:
            bg_type = "white"
        else:
            bg_type = "complex"

        # ── Step 13: Translate ────────────────────────────────
        translated        = ""
        translation_input = corrected_text

        if region_type not in {"sfx"} and is_meaningful_text(corrected_text, source_lang):
            try:
                raw_translation = translator.translate(translation_input)
                translated      = postprocess_translation_general(raw_translation or "")
            except Exception as e:
                print(f"   ⚠️  Translation failed for BOX#{bid}: {e}")
                translated = corrected_text

        # Segment bubble_groups into || separated string for output
        bubble_groups_str = " || ".join(bubble_groups) if bubble_groups else corrected_text

        # Determine OCR source label
        ocr_source = "vision-base"
        if correction_gain > 0.05:
            ocr_source = "vision-reread"

        # Add BUBBLE / SEGMENTED flags
        if bubble_groups and len(bubble_groups) > 1:
            if "BUBBLE" not in flags:
                flags.append("BUBBLE")
            if "SEGMENTED" not in flags:
                flags.append("SEGMENTED")

        results[str(bid)] = {
            "order":             order_idx,
            "region_type":       region_type,
            "background_type":   bg_type,              # <--- NEW FLAG ADDED HERE
            "confidence":        round(conf, 4),
            "ocr_source":        ocr_source,
            "raw_ocr":           raw_text,
            "corrected_ocr":     corrected_text,
            "translation_input": translation_input,
            "translated":        translated,
            "flags":             flags,
            "bubble_groups":     bubble_groups,
            "box":               xyxy_to_xywh(adjusted_box_xyxy),
            "lines":             bubble_groups,
        }

    print(f"\n   ✅ Processed {len(results)} text region(s).")

    # ── Step 14: Write outputs ────────────────────────────────
    if output_json:
        _write_json_output(results, output_json)

    if output_txt:
        _write_txt_output(results, output_txt)

    return results

# ============================================================
# OUTPUT WRITERS
# ============================================================
def _write_json_output(results: Dict[str, Any], path: str) -> None:
    try:
        with open(path, "w", encoding="utf-8") as f:
            json.dump(results, f, ensure_ascii=False, indent=2)
        print(f"   💾 JSON saved → {path}")
    except Exception as e:
        print(f"   ⚠️  Failed to write JSON: {e}")


def _write_txt_output(results: Dict[str, Any], path: str) -> None:
    sep   = "─" * 120
    lines = [
        "BUBBLE|ORDER|TYPE|CONF|OCR_SOURCE|ORIGINAL|CORRECTED|BUBBLE_GROUPS|TRANSLATED|FLAGS",
        sep,
    ]

    for bid, data in sorted(results.items(), key=lambda kv: kv[1]["order"]):
        bubble_groups_str = " || ".join(data.get("bubble_groups", []))
        flags_str         = ",".join(data.get("flags", []))
        row = (
            f"#{bid}"
            f"|{data['order']}"
            f"|{data['region_type']}"
            f"|{data['confidence']:.2f}"
            f"|{data['ocr_source']}"
            f"|{data['raw_ocr']}"
            f"|{data['corrected_ocr']}"
            f"|{bubble_groups_str}"
            f"|{data['translated']}"
            f"|{flags_str}"
        )
        lines.append(row)

    try:
        with open(path, "w", encoding="utf-8") as f:
            f.write("\n".join(lines) + "\n")
        print(f"   📄 TXT saved  → {path}")
    except Exception as e:
        print(f"   ⚠️  Failed to write TXT: {e}")


# ============================================================
# DEBUG VISUALISER
# ============================================================
def draw_debug_clusters(image_bgr: np.ndarray,
                         out_boxes:   Dict[int, tuple],
                         out_lines:   Dict[int, list],
                         out_indices: Dict[int, list],
                         ocr:         list,
                         save_path:   str = None) -> np.ndarray:
    vis    = image_bgr.copy()
    ih, iw = vis.shape[:2]

    COLOR_MAP = {
        "dialogue":  (0,   200,   0),
        "narration": (0,   165, 255),
        "reaction":  (255, 200,   0),
        "sfx":       (0,     0, 220),
        "unknown":   (120, 120, 120),
    }

    bid_order = sorted(
        out_boxes.keys(),
        key=lambda b: (
            (out_boxes[b][1] + out_boxes[b][3]) / 2.0,
            (out_boxes[b][0] + out_boxes[b][2]) / 2.0,
        )
    )

    for order_idx, bid in enumerate(bid_order, start=1):
        box   = out_boxes[bid]
        lines = out_lines.get(bid, [])
        text  = normalize_text(" ".join(lines))

        rtype = region_text_role_hint(text)
        color = COLOR_MAP.get(rtype, (120, 120, 120))

        x1, y1, x2, y2 = box
        cv2.rectangle(vis, (x1, y1), (x2, y2), color, 2)

        label      = f"BOX#{bid} [{rtype}]"
        preview    = (text[:40] + "...") if len(text) > 40 else text
        font       = cv2.FONT_HERSHEY_SIMPLEX
        font_scale = 0.38
        thickness  = 1

        (lw, lh), _ = cv2.getTextSize(label, font, font_scale, thickness)
        cv2.rectangle(vis,
                      (x1, max(0, y1 - lh - 6)),
                      (x1 + lw + 4, y1),
                      color, -1)
        cv2.putText(vis, label,
                    (x1 + 2, max(lh, y1 - 3)),
                    font, font_scale, (255, 255, 255), thickness,
                    cv2.LINE_AA)

        cv2.putText(vis, preview,
                    (x1 + 2, min(ih - 5, y1 + lh + 6)),
                    font, font_scale * 0.85, color, thickness,
                    cv2.LINE_AA)

        # Draw individual OCR quad outlines
        for idx in out_indices.get(bid, []):
            if idx < len(ocr):
                q   = ocr[idx][0]
                pts = np.array(q, dtype=np.int32).reshape((-1, 1, 2))
                cv2.polylines(vis, [pts], True,
                              tuple(min(255, c + 80) for c in color), 1)

    if save_path:
        cv2.imwrite(save_path, vis)
        print(f"   🖼️  Debug image saved → {save_path}")

    return vis


# ============================================================
# CLI ENTRY POINT
# ============================================================
def main():
    import argparse

    parser = argparse.ArgumentParser(
        description="Manga page OCR + translation pipeline (macOS Vision)")
    parser.add_argument("image",           help="Path to manga page image")
    parser.add_argument("--source", "-s",  default="en",
                        help="Source language code (default: en)")
    parser.add_argument("--target", "-t",  default="ca",
                        help="Target language code (default: ca)")
    parser.add_argument("--json",   "-j",  default=None,
                        help="Output JSON file path")
    parser.add_argument("--txt",    "-o",  default=None,
                        help="Output TXT file path")
    parser.add_argument("--debug",  "-d",  default=None,
                        help="Save debug visualisation to this path")
    args = parser.parse_args()

    base      = os.path.splitext(args.image)[0]
    json_out  = args.json  or f"{base}_bubbles.json"
    txt_out   = args.txt   or f"{base}_output.txt"
    debug_out = args.debug or f"{base}_debug_clusters.png"

    results = process_manga_page(
        image_path  = args.image,
        source_lang = args.source,
        target_lang = args.target,
        output_json = json_out,
        output_txt  = txt_out,
    )

    if not results:
        print("\n❌ No results produced.")
        return

    # ── Debug visualisation ───────────────────────────────────
    image_bgr = cv2.imread(args.image)
    if image_bgr is not None:
        vis_boxes:   Dict[int, tuple] = {}
        vis_lines:   Dict[int, list]  = {}
        vis_indices: Dict[int, list]  = {}

        for bid_str, data in results.items():
            bid  = int(bid_str)
            xywh = data["box"]
            vis_boxes[bid] = (
                xywh["x"],
                xywh["y"],
                xywh["x"] + xywh["w"],
                xywh["y"] + xywh["h"],
            )
            vis_lines[bid]   = data.get("lines", [])
            vis_indices[bid] = []

        draw_debug_clusters(
            image_bgr,
            vis_boxes,
            vis_lines,
            vis_indices,
            ocr       = [],
            save_path = debug_out,
        )

    # ── Console summary ───────────────────────────────────────
    print(f"\n{'='*60}")
    print(f"📊 SUMMARY  ({len(results)} boxes)")
    print(f"{'='*60}")
    for bid_str, data in sorted(results.items(), key=lambda kv: kv[1]["order"]):
        print(
            f"  #{bid_str:>3}  [{data['region_type']:<9}]  "
            f"conf={data['confidence']:.2f}  "
            f"\"{data['corrected_ocr'][:55]}\""
        )
    print(f"{'='*60}\n")


if __name__ == "__main__":
    main()