"""签名图片处理诊断脚本（独立调试用）。

工作流:
    把待测试的图片 (.png / .jpg / .jpeg) 丢到 image_test/input/ 下，
    直接双击 run.bat 或 `python image_test/process.py` 即可。
    脚本会自动遍历 input/ 下所有图片，每张图都跑完整流水线并把
    每一步中间结果存到 image_test/output/<图片名>/ 子目录。

参数调整: 直接改下方 CONFIG 段，不用命令行。

输出结构 (每张图一个子目录):
    image_test/output/<stem>/
        00_original.png         原图
        01_binarized.png        二值化后（仅 is_signature=true 时）
        02_white_removed.png    白底转透明后
        03_border_cropped.png   透明边框裁剪后
        04_scaled.png           按高度等比缩放后
        05_noised.png           边缘噪声后
        06_rotated.png          旋转后
        07_final_cropped.png    紧致裁剪后（最终输出）
        final.png               最终输出（等价于插入 Word 的那张）
        diagnostics.txt         每步尺寸 / 内容像素占比 / bbox
"""
import io
import random
import shutil
from pathlib import Path

from PIL import Image, ImageChops

# ============================================================================
# 配置区：直接改这里，不用命令行
# ============================================================================
CONFIG = {
    # 处理对象
    "is_signature": True,        # True: 手写签字(先二值化)；False: 普通签章(不二值化)

    # 处理参数（对应正式流水线的全局参数）
    "height_cm": 1,            # 目标高度(cm)；None 表示不缩放
    "noise_level": 0,            # 边缘噪声 0~10
    "rotate_min": 0,             # 旋转最小角度
    "rotate_max": 0,             # 旋转最大角度
    "dpi": 96,                   # 缩放 DPI

    # 算法阈值（出问题时优先调这几个）
    "alpha_threshold": 30,       # 紧致裁剪: alpha 高于此值视为有内容
    "brightness_threshold": 210, # 紧致裁剪: 亮度低于此值视为暗（笔画）
    "binarize_threshold": 128,   # 二值化阈值
    "white_threshold": 240,      # 抠白底: RGB 都高于此值视为白底

    # 调试控制
    "stop_at": "all",            # all / original / binarize / remove_white / crop_border / scale / noise / rotate
    "seed": 42,                  # 随机种子（让噪声/旋转结果可复现；None=随机）
    "clear_output": True,        # 每次运行前清空 output 目录
}

# 输入输出目录（相对脚本位置，固定）
BASE_DIR = Path(__file__).resolve().parent
INPUT_DIR = BASE_DIR / "input"
OUTPUT_DIR = BASE_DIR / "output"

IMAGE_EXTS = {".png", ".jpg", ".jpeg"}


# ============================================================================
# 以下函数从 lib/image_processor.py 原样复制，方便独立调参
# ============================================================================

def binarize(img, threshold=128):
    """手写签字黑白二值化：>threshold -> 255(白)，否则 0(黑)。"""
    if img.mode != 'L':
        img = img.convert('L')
    img = img.point(lambda p: 255 if p > threshold else 0)
    img = img.convert('RGBA')
    return img


def remove_white_background(img, white_threshold=240):
    """接近白色的像素转为全透明。"""
    if img.mode != 'RGBA':
        img = img.convert('RGBA')
    datas = img.getdata()
    new_data = []
    for item in datas:
        r, g, b, a = item
        if r > white_threshold and g > white_threshold and b > white_threshold:
            new_data.append((255, 255, 255, 0))
        else:
            new_data.append((r, g, b, a))
    img.putdata(new_data)
    return img


def crop_transparent_border(img):
    """按 alpha 通道 getbbox 裁掉四周完全透明区域。"""
    if img.mode != 'RGBA':
        img = img.convert('RGBA')
    bbox = img.getbbox()
    if bbox:
        img = img.crop(bbox)
    return img


def scale_by_height(img, height_cm, dpi=96):
    """按目标高度（厘米）等比缩放。"""
    try:
        from PIL.Image import Resampling
        resample = Resampling.LANCZOS
    except ImportError:
        resample = Image.LANCZOS
    target_h = int(height_cm * dpi / 2.54)
    if target_h <= 0:
        return img
    w, h = img.size
    if h == 0:
        return img
    ratio = target_h / h
    new_w = max(1, int(w * ratio))
    new_h = target_h
    return img.resize((new_w, new_h), resample)


def add_edge_noise(img, noise_level):
    """在笔画边缘添加随机噪点，模拟印章不均匀。"""
    if img.mode != 'RGBA':
        img = img.convert('RGBA')
    w, h = img.size
    pixels = img.load()
    intensity = noise_level * 2
    for x in range(w):
        for y in range(h):
            r, g, b, a = pixels[x, y]
            if a == 0:
                continue
            is_edge = False
            for dx in range(-2, 3):
                for dy in range(-2, 3):
                    nx, ny = x + dx, y + dy
                    if 0 <= nx < w and 0 <= ny < h:
                        if pixels[nx, ny][3] == 0:
                            is_edge = True
                            break
                if is_edge:
                    break
            if is_edge and random.random() < noise_level / 10.0:
                offset_x = random.randint(-intensity, intensity)
                offset_y = random.randint(-intensity, intensity)
                nx = x + offset_x
                ny = y + offset_y
                if 0 <= nx < w and 0 <= ny < h:
                    nr = min(255, max(0, r + random.randint(-20, 20)))
                    ng = min(255, max(0, g + random.randint(-20, 20)))
                    nb = min(255, max(0, b + random.randint(-20, 20)))
                    na = min(255, max(0, a + random.randint(-30, 10)))
                    if pixels[nx, ny][3] == 0:
                        pixels[nx, ny] = (nr, ng, nb, na)
    return img


def rotate_image(img, rotate_min, rotate_max):
    """随机角度旋转，expand=True 扩展画布以容纳旋转后内容。"""
    try:
        from PIL.Image import Resampling
        resample = Resampling.BICUBIC
    except ImportError:
        resample = Image.BICUBIC
    if rotate_min > rotate_max:
        rotate_min, rotate_max = rotate_max, rotate_min
    angle = random.randint(rotate_min, rotate_max)
    return img.rotate(angle, resample=resample, expand=True,
                      fillcolor=(0, 0, 0, 0))


def crop_to_content(img, alpha_threshold=30, brightness_threshold=210):
    """末尾紧致裁剪：(alpha > th) AND (亮度 < th)。"""
    if img.mode != 'RGBA':
        img = img.convert('RGBA')
    r, g, b, a = img.split()
    luma = Image.merge('RGB', (r, g, b)).convert('L')
    alpha_mask = a.point(lambda p: 255 if p > alpha_threshold else 0)
    luma_mask = luma.point(lambda p: 255 if p < brightness_threshold else 0)
    combined = ImageChops.darker(alpha_mask, luma_mask)
    bbox = combined.getbbox()
    if bbox and bbox != (0, 0, img.size[0], img.size[1]):
        img = img.crop(bbox)
    return img, bbox


# ============================================================================
# 诊断辅助
# ============================================================================

def content_stats(img, alpha_threshold=30, brightness_threshold=210):
    if img.mode != 'RGBA':
        img = img.convert('RGBA')
    r, g, b, a = img.split()
    luma = Image.merge('RGB', (r, g, b)).convert('L')
    alpha_mask = a.point(lambda p: 255 if p > alpha_threshold else 0)
    luma_mask = luma.point(lambda p: 255 if p < brightness_threshold else 0)
    combined = ImageChops.darker(alpha_mask, luma_mask)
    total = img.size[0] * img.size[1] or 1
    ha, hl, hc = alpha_mask.histogram(), luma_mask.histogram(), combined.histogram()
    return {
        'size': img.size,
        'alpha_pct': 100.0 * sum(ha[1:]) / total,
        'dark_pct': 100.0 * sum(hl[1:]) / total,
        'content_pct': 100.0 * sum(hc[1:]) / total,
        'bbox': combined.getbbox(),
    }


def save_step(out_dir, idx, name, img, diag_lines, alpha_th, bright_th):
    if img is None:
        diag_lines.append(f"[{idx:02d}] {name:18s} SKIPPED")
        return
    path = out_dir / f"{idx:02d}_{name}.png"
    img.save(path)
    s = content_stats(img, alpha_th, bright_th)
    line = (f"[{idx:02d}] {name:18s} size={str(s['size']):14s} "
            f"alpha={s['alpha_pct']:5.1f}%  dark={s['dark_pct']:5.1f}%  "
            f"content={s['content_pct']:5.1f}%  bbox={s['bbox']}")
    diag_lines.append(line)


# ============================================================================
# 单张图处理
# ============================================================================

def process_one(input_path, out_dir, cfg):
    out_dir = Path(out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    diag = []
    img = Image.open(input_path)
    diag.append(f"输入: {input_path.name}")
    diag.append(f"模式: {img.mode}  尺寸: {img.size}")
    diag.append(f"参数: is_sig={cfg['is_signature']} height_cm={cfg['height_cm']} "
                f"noise={cfg['noise_level']} rotate=[{cfg['rotate_min']},{cfg['rotate_max']}] dpi={cfg['dpi']}")
    diag.append(f"阈值: alpha>{cfg['alpha_threshold']}  亮度<{cfg['brightness_threshold']}  "
                f"二值化>{cfg['binarize_threshold']}  白底>{cfg['white_threshold']}")
    diag.append("-" * 80)

    save_step(out_dir, 0, "original", img, diag,
              cfg['alpha_threshold'], cfg['brightness_threshold'])
    if cfg['stop_at'] == 'original':
        return diag

    if cfg['is_signature']:
        img = binarize(img, threshold=cfg['binarize_threshold'])
    save_step(out_dir, 1, "binarized",
              img if cfg['is_signature'] else None, diag,
              cfg['alpha_threshold'], cfg['brightness_threshold'])
    if cfg['stop_at'] == 'binarize':
        return diag

    img = remove_white_background(img, white_threshold=cfg['white_threshold'])
    save_step(out_dir, 2, "white_removed", img, diag,
              cfg['alpha_threshold'], cfg['brightness_threshold'])
    if cfg['stop_at'] == 'remove_white':
        return diag

    img = crop_transparent_border(img)
    save_step(out_dir, 3, "border_cropped", img, diag,
              cfg['alpha_threshold'], cfg['brightness_threshold'])
    if cfg['stop_at'] == 'crop_border':
        return diag

    h = cfg['height_cm']
    if h is not None and h > 0:
        img = scale_by_height(img, h, dpi=cfg['dpi'])
    save_step(out_dir, 4, "scaled", img if (h is not None and h > 0) else None,
              diag, cfg['alpha_threshold'], cfg['brightness_threshold'])
    if cfg['stop_at'] == 'scale':
        return diag

    if cfg['noise_level'] > 0:
        img = add_edge_noise(img, cfg['noise_level'])
    save_step(out_dir, 5, "noised", img if cfg['noise_level'] > 0 else None,
              diag, cfg['alpha_threshold'], cfg['brightness_threshold'])
    if cfg['stop_at'] == 'noise':
        return diag

    if cfg['rotate_min'] != 0 or cfg['rotate_max'] != 0:
        img = rotate_image(img, cfg['rotate_min'], cfg['rotate_max'])
    save_step(out_dir, 6, "rotated",
              img if (cfg['rotate_min'] != 0 or cfg['rotate_max'] != 0) else None,
              diag, cfg['alpha_threshold'], cfg['brightness_threshold'])
    if cfg['stop_at'] == 'rotate':
        return diag

    img, _ = crop_to_content(img,
                             alpha_threshold=cfg['alpha_threshold'],
                             brightness_threshold=cfg['brightness_threshold'])
    save_step(out_dir, 7, "final_cropped", img, diag,
              cfg['alpha_threshold'], cfg['brightness_threshold'])

    final_buf = io.BytesIO()
    img.save(final_buf, format='PNG')
    with open(out_dir / "final.png", 'wb') as f:
        f.write(final_buf.getvalue())
    diag.append("-" * 80)
    diag.append(f"最终输出: final.png  size={img.size}")
    return diag


# ============================================================================
# 主入口：扫描 input/ 下所有图片
# ============================================================================

def main():
    if not INPUT_DIR.exists():
        INPUT_DIR.mkdir(parents=True, exist_ok=True)

    if CONFIG.get('seed') is not None:
        random.seed(CONFIG['seed'])

    images = sorted([p for p in INPUT_DIR.iterdir()
                     if p.suffix.lower() in IMAGE_EXTS and p.is_file()])

    if not images:
        print(f"未在 {INPUT_DIR} 下找到 png/jpg 图片，请先把测试图片放进去")
        return

    if CONFIG.get('clear_output') and OUTPUT_DIR.exists():
        shutil.rmtree(OUTPUT_DIR)
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

    print(f"找到 {len(images)} 张图: {[p.name for p in images]}")
    print(f"输出目录: {OUTPUT_DIR}\n")

    for img_path in images:
        out_sub = OUTPUT_DIR / img_path.stem
        print(f"=== 处理: {img_path.name} ===")
        try:
            diag = process_one(img_path, out_sub, CONFIG)
            for line in diag:
                print(line)
            with open(out_sub / "diagnostics.txt", 'w', encoding='utf-8') as f:
                f.write("\n".join(diag))
            print(f"-> 输出: {out_sub}\n")
        except Exception as e:
            print(f"处理失败: {e}\n")


if __name__ == '__main__':
    main()