#!/usr/bin/env python3
"""Split long RTL OCR lines into CTC-safe Persian word-group crops."""

from __future__ import annotations

import argparse
from pathlib import Path

import cv2
import numpy as np
from PIL import Image


def zero_runs(mask: np.ndarray) -> list[tuple[int, int]]:
    runs = []
    start = None
    for index, value in enumerate(mask.tolist() + [False]):
        if value and start is None:
            start = index
        elif not value and start is not None:
            runs.append((start, index))
            start = None
    return runs


def word_intervals(image: Image.Image, count: int) -> list[tuple[int, int]] | None:
    if count == 1:
        return [(0, image.width)]
    gray = np.asarray(image.convert("L"))
    _, ink = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
    projection = (ink > 0).sum(axis=0)
    active = np.flatnonzero(projection > max(1, image.height // 80))
    if active.size == 0:
        return None
    left, right = int(active[0]), int(active[-1]) + 1
    gaps = [
        (a, b) for a, b in zero_runs(projection[left:right] <= max(1, image.height // 100))
        if b - a >= 2 and a > 0 and b < right - left
    ]
    if len(gaps) < count - 1:
        return None
    chosen = sorted(
        ((a + left, b + left) for a, b in sorted(gaps, key=lambda run: run[1] - run[0], reverse=True)[: count - 1]),
        key=lambda run: run[0],
    )
    cuts = [left] + [(a + b) // 2 for a, b in chosen] + [right]
    intervals = [(cuts[i], cuts[i + 1]) for i in range(len(cuts) - 1)]
    return list(reversed(intervals))


def groups(words: list[str], intervals: list[tuple[int, int]], max_chars: int):
    current_words: list[str] = []
    current_boxes: list[tuple[int, int]] = []
    for word, box in zip(words, intervals):
        candidate = " ".join(current_words + [word])
        if current_words and len(candidate) > max_chars:
            yield " ".join(current_words), current_boxes
            current_words, current_boxes = [], []
        if len(word) > max_chars:
            return
        current_words.append(word)
        current_boxes.append(box)
    if current_words:
        yield " ".join(current_words), current_boxes


def process_list(source: Path, destination: Path, image_dir: Path, root: Path, max_chars: int) -> tuple[int, int]:
    output = []
    rejected = 0
    for line_index, raw in enumerate(source.read_text(encoding="utf-8").splitlines()):
        relative, text = raw.split("\t", 1)
        words = text.split()
        with Image.open(root / relative) as opened:
            image = opened.convert("RGB")
            intervals = word_intervals(image, len(words))
            if intervals is None:
                rejected += 1
                continue
            for group_index, (label, boxes) in enumerate(groups(words, intervals, max_chars)):
                x0 = max(0, min(x[0] for x in boxes) - 3)
                x1 = min(image.width, max(x[1] for x in boxes) + 3)
                crop = image.crop((x0, 0, x1, image.height))
                path = image_dir / f"{line_index:07d}-{group_index:03d}.png"
                crop.save(path)
                output.append(f"{path.relative_to(root)}\t{label}")
    destination.write_text("\n".join(output) + "\n", encoding="utf-8")
    return len(output), rejected


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--data-root", type=Path, required=True)
    parser.add_argument("--pilot-dir", type=Path, required=True)
    parser.add_argument("--max-chars", type=int, default=25)
    args = parser.parse_args()
    for split in ("train", "val"):
        image_dir = args.pilot_dir / "images" / f"segments-{split}"
        image_dir.mkdir(parents=True, exist_ok=True)
        kept, rejected = process_list(
            args.pilot_dir / f"{split}_list.txt",
            args.pilot_dir / f"{split}_segments.txt",
            image_dir,
            args.data_root,
            args.max_chars,
        )
        print(f"{split}: segments={kept} rejected_lines={rejected}")


if __name__ == "__main__":
    main()