"""
提示词图像自动切割模块
将 170x50 的提示图自动切分为 3 个独立的图标小图
"""

from pathlib import Path
from typing import List, Tuple
import numpy as np
from PIL import Image


import io

def slice_prompt_image(prompt_img: str | Path | bytes | Image.Image, padding: int = 2) -> List[Tuple[Image.Image, Tuple[int, int, int, int]]]:
    """
    自动识别并切割提示图中的 3 个字符/图标
    返回: [(PIL_Image, (x1, y1, x2, y2)), ...]
    """
    if isinstance(prompt_img, Image.Image):
        img = prompt_img.convert("L")
    elif isinstance(prompt_img, bytes):
        img = Image.open(io.BytesIO(prompt_img)).convert("L")
    else:
        img = Image.open(prompt_img).convert("L")
    arr = np.array(img)
    h, w = arr.shape

    # 背景灰度基准 (通常取左上角像素)
    bg_val = int(arr[0, 0])

    # 提取暗色笔画像素（排除灰色背景）
    mask = arr < (bg_val - 25)

    # 水平方向投影，找出连续包含笔画的列区间
    cols_has_dark = mask.any(axis=0)

    segments: List[Tuple[int, int]] = []
    in_seg = False
    start = 0

    for i, val in enumerate(cols_has_dark):
        if val and not in_seg:
            in_seg = True
            start = i
        elif not val and in_seg:
            in_seg = False
            if i - start >= 5:  # 过滤过窄杂点
                segments.append((start, i))
    if in_seg and (w - start >= 5):
        segments.append((start, w))

    # 若投影分割不足 3 段（极罕见），按固定 3 等分兜底
    if len(segments) != 3:
        step = w // 3
        segments = [(i * step + 5, (i + 1) * step - 5) for i in range(3)]

    # 原始彩色图用于切片
    if isinstance(prompt_img, Image.Image):
        orig_img = prompt_img.convert("RGBA")
    elif isinstance(prompt_img, bytes):
        orig_img = Image.open(io.BytesIO(prompt_img)).convert("RGBA")
    else:
        orig_img = Image.open(prompt_img).convert("RGBA")
    crops = []

    for seg_start, seg_end in segments[:3]:
        # 纵向精确裁剪
        sub_mask = mask[:, seg_start:seg_end]
        rows_has_dark = sub_mask.any(axis=1)
        row_indices = np.where(rows_has_dark)[0]

        if len(row_indices) > 0:
            top = max(0, int(row_indices[0]) - padding)
            bottom = min(h, int(row_indices[-1]) + 1 + padding)
        else:
            top, bottom = 0, h

        left = max(0, seg_start - padding)
        right = min(w, seg_end + padding)

        box = (left, top, right, bottom)
        crop = orig_img.crop(box)
        crops.append((crop, box))

    return crops


if __name__ == "__main__":
    import sys

    sample = sys.argv[1] if len(sys.argv) > 1 else "tcaptcha/data/raw/1789135559041_4170_prompt.png"
    if Path(sample).exists():
        results = slice_prompt_image(sample)
        print(f"切片完成，切出 {len(results)} 个目标:")
        for idx, (c, b) in enumerate(results):
            print(f"  [{idx+1}] 尺寸: {c.size}, 边界框: {b}")
    else:
        print(f"样本文件不存在: {sample}")