from __future__ import annotations

import argparse
import json
from pathlib import Path
import uuid
import zlib
import cv2
import matplotlib.pyplot as plt
import numpy as np
import tifffile
from reedsolo import RSCodec, ReedSolomonError
from tardos import UserRecord, accuse_symmetric, build_codebook, rank_suspects, sample_tardos_probabilities
from watermark import BlindDwtDctQim, average_collusion, xor_collusion_via_reembed
from scan_utils import detect_page_quad
from psb_bridge import (
    SourceMeta,
    read_cover_asset,
    save_bgr_preview_png,
    save_bgr_as_flattened_psb_via_photoshop,
    save_bgr_native_tiff,
    prepare_large_raster_input,
    clone_raster_file,
    memmap_tiff_rgb,
    release_memmap,
    save_tiff_preview_png,
    save_flattened_tiff_as_psb_via_photoshop,
)


DISPLAY_TEXT_BYTES = 40  # 可读字符串正文最大 40 字节
DISPLAY_MAGIC = b"WM"
DISPLAY_VERSION = 1
DISPLAY_RS_NSYM = 16  # RS 冗余字节数，可先从16开始

_DISPLAY_PACKET_BYTES = 2 + 1 + 1 + DISPLAY_TEXT_BYTES + 4
_DISPLAY_CODEWORD_BYTES = _DISPLAY_PACKET_BYTES + DISPLAY_RS_NSYM
DISPLAY_BITS_LEN = _DISPLAY_CODEWORD_BYTES * 8

_RS = RSCodec(DISPLAY_RS_NSYM)

TEXT_SUFFIX = "::text"
FP_SUFFIX = "::fp"


def bytes_to_bits(data: bytes) -> np.ndarray:
    return np.unpackbits(np.frombuffer(data, dtype=np.uint8)).astype(np.uint8)


def bits_to_bytes(bits: np.ndarray) -> bytes:
    bits = np.asarray(bits, dtype=np.uint8).flatten()
    pad = (-len(bits)) % 8
    if pad:
        bits = np.pad(bits, (0, pad), constant_values=0)
    return np.packbits(bits).tobytes()


def build_display_text(user_id: str, content_id: str) -> str:
    return f"UID={user_id};CID={content_id}"


def _rs_decode_to_bytes(decoded_obj) -> bytes:
    if isinstance(decoded_obj, tuple):
        return bytes(decoded_obj[0])
    return bytes(decoded_obj)


def encode_display_text(text: str, max_text_bytes: int = DISPLAY_TEXT_BYTES) -> np.ndarray:
    raw = text.encode("utf-8")
    if len(raw) > max_text_bytes:
        raise ValueError(f"展示字符串过长，最多 {max_text_bytes} 字节，当前 {len(raw)}")

    body = raw.ljust(max_text_bytes, b"\0")
    crc = zlib.crc32(raw).to_bytes(4, "big")

    packet = (
        DISPLAY_MAGIC +
        bytes([DISPLAY_VERSION]) +
        bytes([len(raw)]) +
        body +
        crc
    )

    codeword = _RS.encode(packet)
    return bytes_to_bits(bytes(codeword))


def decode_display_text(bits: np.ndarray, max_text_bytes: int = DISPLAY_TEXT_BYTES) -> tuple[str, bool]:
    data = bits_to_bytes(bits)
    if len(data) < _DISPLAY_CODEWORD_BYTES:
        return "", False

    codeword = data[:_DISPLAY_CODEWORD_BYTES]

    try:
        packet = _rs_decode_to_bytes(_RS.decode(codeword))
    except ReedSolomonError:
        return "", False
    except Exception:
        return "", False

    if len(packet) < _DISPLAY_PACKET_BYTES:
        return "", False

    if packet[:2] != DISPLAY_MAGIC:
        return "", False

    version = packet[2]
    if version != DISPLAY_VERSION:
        return "", False

    n = int(packet[3])
    if n > max_text_bytes:
        return "", False

    raw = packet[4:4 + n]
    crc_saved = packet[4 + max_text_bytes: 4 + max_text_bytes + 4]
    crc_now = zlib.crc32(raw).to_bytes(4, "big")
    if crc_saved != crc_now:
        return "", False

    try:
        text = raw.decode("utf-8")
    except Exception:
        return "", False

    return text, True

USER_DB = [
    UserRecord(user_id="user_1",  user_key="k_9f2A1x7P"),
    UserRecord(user_id="user_2",  user_key="k_B8mQ3d2L"),
    UserRecord(user_id="user_3",  user_key="k_T4zN9v1R"),
    UserRecord(user_id="user_4",  user_key="k_H7pX5s8W"),
    UserRecord(user_id="user_5",  user_key="k_M2cJ6u4K"),
    UserRecord(user_id="user_6",  user_key="k_R5yD8n3F"),
    UserRecord(user_id="user_7",  user_key="k_V1qL7b9T"),
    UserRecord(user_id="user_8",  user_key="k_P6wE2m5C"),
    UserRecord(user_id="user_9",  user_key="k_N3rK4x8Z"),
    UserRecord(user_id="user_10", user_key="k_G8tU1p6Y"),
]

def make_synthetic_cover(size: int = 512) -> np.ndarray:
    y, x = np.mgrid[0:size, 0:size].astype(np.float32)
    base = 110 + 40 * np.sin(x / 17.0) + 35 * np.cos(y / 29.0)
    grad = 0.25 * x + 0.15 * y
    tex = 12 * np.sin((x + y) / 11.0) + 8 * np.cos((x - y) / 19.0)
    b = np.clip(base + grad * 0.15 + tex, 0, 255)
    g = np.clip(base + grad * 0.10 - tex * 0.8 + 20, 0, 255)
    r = np.clip(base + grad * 0.12 + tex * 0.5 + 35, 0, 255)
    img = np.stack([b, g, r], axis=-1).astype(np.uint8)
    cv2.circle(img, (size // 2, size // 2), size // 5, (220, 160, 90), thickness=-1)
    cv2.rectangle(img, (40, 60), (190, 170), (80, 200, 230), thickness=-1)
    cv2.putText(img, "WM DEMO", (size // 5, size - 55),
                cv2.FONT_HERSHEY_SIMPLEX, 1.8, (245, 245, 245), 3, cv2.LINE_AA)
    return img

def dtype_peak(dtype: np.dtype) -> float:
    dtype = np.dtype(dtype)
    if dtype == np.uint8:
        return 255.0
    if dtype == np.uint16:
        return 65535.0
    raise ValueError(f"不支持的 dtype: {dtype}")


def to_preview_u8(image: np.ndarray) -> np.ndarray:
    if image.dtype == np.uint8:
        return image

    if image.dtype == np.uint16:
        preview = np.rint(image.astype(np.float32) / 65535.0 * 255.0)
        return np.clip(preview, 0, 255).astype(np.uint8)

    raise ValueError(f"不支持预览转换的 dtype: {image.dtype}")

def save_rgb(path: Path, bgr: np.ndarray) -> None:
    cv2.imwrite(str(path), to_preview_u8(bgr))


def plot_scores(path: Path, score_dict: dict[str, list[tuple[str, float]]], colluders: set[str]) -> None:
    fig, axes = plt.subplots(1, len(score_dict), figsize=(6 * len(score_dict), 5))
    if len(score_dict) == 1:
        axes = [axes]

    for ax, (title, ranking) in zip(axes, score_dict.items()):
        users = [u for u, _ in ranking]
        scores = [s for _, s in ranking]
        colors = ["tab:red" if u in colluders else "tab:blue" for u in users]
        ax.bar(users, scores, color=colors)
        ax.set_title(title)
        ax.set_ylabel("Tardos score")
        ax.tick_params(axis="x", rotation=45)

    fig.tight_layout()
    fig.savefig(path, dpi=180)
    plt.close(fig)


def parse_args():
    parser = argparse.ArgumentParser(description="Blind watermark demo with Tardos + DWT-DCT + QIM")

    parser.add_argument("--image", type=str, default=None,
                        help="待嵌入水印的图片路径；不传则自动生成合成图")
    parser.add_argument("--out-dir", type=str, default="results",
                        help="输出目录")

    parser.add_argument("--num-users", type=int, default=8,
                        help="用户数量")
    parser.add_argument("--code-length", type=int, default=320,
                        help="Tardos码长")
    parser.add_argument("--colluders", type=str, default="user_2,user_5,user_7",
                        help="合谋用户，逗号分隔，例如 user_2,user_5,user_7")

    parser.add_argument("--delta", type=float, default=8.0,
                        help="QIM量化步长")
    parser.add_argument("--repeats", type=int, default=2,
                        help="每个bit重复嵌入次数")
    parser.add_argument("--delta-mode", type=str, default="normalized_8bit",
                    choices=["normalized_8bit", "native"],
                    help="delta 的解释方式")
    
    parser.add_argument("--master-key", type=str, default="server-master-key-2026",
                        help="系统主密钥")
    parser.add_argument("--content-id", type=str, default="asset-demo-001",
                        help="内容ID")
    parser.add_argument("--seed", type=int, default=2026,
                        help="Tardos概率向量随机种子")

    parser.add_argument("--cover-size", type=int, default=512,
                        help="未提供图片时，自动生成合成图的尺寸")
    parser.add_argument("--emit-psb", action="store_true",
                        help="如果输入是 PSD/PSB，则额外输出扁平单层 PSB 分发文件")

    parser.add_argument("--psb-max-side", type=int, default=0,
                        help="读取 PSD/PSB 时可选缩边；0 表示不缩放，直接使用原始尺寸（可能导致内存占用过大）")
    return parser.parse_args()


# def load_cover_image(image_path: str | None, cover_size: int) -> np.ndarray:
#     if image_path is None:
#         return make_synthetic_cover(cover_size)

#     img = cv2.imread(image_path, cv2.IMREAD_COLOR)
#     if img is None:
#         raise FileNotFoundError(f"无法读取图片: {image_path}")
#     return img

def load_cover_and_meta(
    image_path: str | None,
    cover_size: int,
    psb_max_side: int,
) -> tuple[np.ndarray, SourceMeta]:
    if image_path is None:
        img = make_synthetic_cover(cover_size)
        meta = SourceMeta(
            input_path=None,
            input_ext=".png",
            input_is_psb=False,
            width=int(img.shape[1]),
            height=int(img.shape[0]),
            depth=8,
        )
        return img, meta

    return read_cover_asset(image_path, max_side=psb_max_side)

def save_main_asset(
    out_dir: Path,
    stem: str,
    bgr: np.ndarray,
    source_meta: SourceMeta,
    emit_psb: bool,
) -> dict:
    result = {
        "png": None,
        "tif": None,
        "psb": None,
        "psb_error": None,
    }
    # 1) 正式保真输出：TIFF
    tif_path = out_dir / f"{stem}.tif"
    save_bgr_native_tiff(tif_path, bgr, source_meta)
    result["tif"] = str(tif_path)

    # 2) 网页/演示预览：PNG
    png_path = out_dir / f"{stem}.png"
    save_bgr_preview_png(png_path, bgr)
    result["png"] = str(png_path)

    # 3) PSB 分发输出
    if emit_psb:
        psb_path = out_dir / f"{stem}.psb"
        try:
            save_bgr_as_flattened_psb_via_photoshop(psb_path, bgr, dpi_x=source_meta.dpi_x, dpi_y=source_meta.dpi_y, icc_profile=source_meta.icc_profile)
            result["psb"] = str(psb_path)
        except Exception as exc:
            result["psb_error"] = str(exc)

    return result

def build_users(num_users: int) -> list[UserRecord]:
    if num_users < 1:
        raise ValueError("num_users 必须 >= 1")
    if num_users > len(USER_DB):
        raise ValueError(f"当前代码里只维护了 {len(USER_DB)} 个用户，请把 num_users 设为 <= {len(USER_DB)}")
    return USER_DB[:num_users]


def parse_colluders(colluders_str: str, valid_user_ids: set[str]) -> list[str]:
    colluders = [x.strip() for x in colluders_str.split(",") if x.strip()]
    if not colluders:
        raise ValueError("colluders 不能为空")

    invalid = [x for x in colluders if x not in valid_user_ids]
    if invalid:
        raise ValueError(f"非法合谋用户ID: {invalid}，可选用户: {sorted(valid_user_ids)}")
    return colluders


def run_demo(args) -> dict:
    out = Path(args.out_dir)
    out.mkdir(parents=True, exist_ok=True)

    # cover = load_cover_image(args.image, args.cover_size)
    # save_rgb(out / "cover.png", cover)

    cover, source_meta = load_cover_and_meta(
        args.image,
        args.cover_size,
        args.psb_max_side,
    )
    if source_meta.input_is_psb and source_meta.depth == 16 and cover.dtype != np.uint16:
        raise RuntimeError(
            f"检测到输入是 16-bit PSB，但当前读入数组 dtype={cover.dtype}。"
            "这说明输入桥没有保住 16-bit，当前结果不能视为保真输出。"
        )
    if source_meta.input_is_psb and args.psb_max_side != 0:
        raise RuntimeError("保真模式下，PSB 输入时 psb_max_side 必须为 0；否则会发生缩图。")
    cover_paths = save_main_asset(out, "cover", cover, source_meta, emit_psb=args.emit_psb)

    users = build_users(args.num_users)
    user_ids = {u.user_id for u in users}
    colluder_ids = parse_colluders(args.colluders, user_ids)

    p = sample_tardos_probabilities(
        length=args.code_length,
        cutoff=0.08,
        seed=args.seed,
    )
    codebook = build_codebook(users, p, master_key=args.master_key)

    wm_text, wm_fp = build_watermarkers(
        master_key=args.master_key,
        delta=args.delta,
        repeats=args.repeats,
        delta_mode=args.delta_mode,
    )

    watermarked = {}
    clean_extract_ber = {}
    psnr_map = {}
    user_outputs = {}
    user_results = {}

    for idx, user in enumerate(users):
        display_text, text_bits, fingerprint_bits = build_single_user_channels(
            user_id=user.user_id,
            user_key=user.user_key,
            content_id=args.content_id,
            master_key=args.master_key,
            code_length=args.code_length,
            seed=args.seed,
        )

        # 先嵌字符串通道
        img_after_text = wm_text.embed(
            cover,
            text_bits,
            content_id=f"{args.content_id}{TEXT_SUFFIX}",
        ).image

        # 再嵌指纹通道
        img = img_after_text
        if len(fingerprint_bits) > 0:
            img = wm_fp.embed(
                img_after_text,
                fingerprint_bits,
                content_id=f"{args.content_id}{FP_SUFFIX}",
            ).image

        watermarked[user.user_id] = img

        # 提字符串通道
        extracted_text_bits = wm_text.extract(
            img,
            n_bits=len(text_bits),
            content_id=f"{args.content_id}{TEXT_SUFFIX}",
        )
        extracted_text, text_ok = decode_display_text(extracted_text_bits)

        # 提指纹通道
        extracted_fingerprint_bits = np.zeros(0, dtype=np.uint8)
        fingerprint_ber = None
        if len(fingerprint_bits) > 0:
            extracted_fingerprint_bits = wm_fp.extract(
                img,
                n_bits=len(fingerprint_bits),
                content_id=f"{args.content_id}{FP_SUFFIX}",
            )
            fingerprint_ber = wm_fp.bit_error_rate(
                fingerprint_bits,
                extracted_fingerprint_bits,
            )

        clean_extract_ber[user.user_id] = fingerprint_ber if fingerprint_ber is not None else 0.0

        peak = dtype_peak(cover.dtype)
        mse = float(np.mean((cover.astype(np.float32) - img.astype(np.float32)) ** 2))
        psnr_map[user.user_id] = 99.0 if mse <= 1e-12 else float(
            10.0 * np.log10((peak ** 2) / mse)
        )

        user_paths = save_main_asset(
            out, f"{user.user_id}_watermarked", img, source_meta, emit_psb=args.emit_psb
        )
        user_outputs[user.user_id] = user_paths

        user_results[user.user_id] = {
            "embedded_text": display_text,
            "extracted_text": extracted_text,
            "text_ok": text_ok,
            "fingerprint_ber": fingerprint_ber,
        }

        if len(extracted_fingerprint_bits) > 0:
            save_rgb(
                out / f"{user.user_id}_extracted_bits.png",
                wm_fp.bits_to_noise_image(extracted_fingerprint_bits)
            )


    colluder_imgs = [watermarked[cid] for cid in colluder_ids]

    pirate_avg = average_collusion(colluder_imgs)
    #total_bits = DISPLAY_BITS_LEN + args.code_length
    text_bits_len = DISPLAY_BITS_LEN
    fp_bits_len = args.code_length

    avg_text_bits = wm_text.extract(
        pirate_avg,
        n_bits=text_bits_len,
        content_id=f"{args.content_id}{TEXT_SUFFIX}",
    )
    avg_text, avg_text_ok = decode_display_text(avg_text_bits)

    avg_fp_bits = np.zeros(fp_bits_len, dtype=np.uint8)
    if fp_bits_len > 0:
        avg_fp_bits = wm_fp.extract(
            pirate_avg,
            n_bits=fp_bits_len,
            content_id=f"{args.content_id}{FP_SUFFIX}",
        )

    avg_scores = accuse_symmetric(avg_fp_bits, codebook, p)
    avg_ranking = rank_suspects(users, avg_scores)

    pirate_avg_paths = save_main_asset(
        out, "pirate_average", pirate_avg, source_meta, emit_psb=args.emit_psb
    )
    if len(avg_fp_bits) > 0:
        save_rgb(out / "pirate_average_bits.png", wm_fp.bits_to_noise_image(avg_fp_bits))

    pirate_xor, xor_bits_fused = xor_collusion_via_reembed(
        base_image=pirate_avg,
        colluder_images=colluder_imgs,
        wm=wm_fp,
        n_bits=fp_bits_len,
        content_id=f"{args.content_id}{FP_SUFFIX}",
    )

    # XOR图上的字符串通道照提
    xor_text_bits = wm_text.extract(
        pirate_xor,
        n_bits=text_bits_len,
        content_id=f"{args.content_id}{TEXT_SUFFIX}",
    )
    xor_text, xor_text_ok = decode_display_text(xor_text_bits)

    xor_fp_bits = np.zeros(fp_bits_len, dtype=np.uint8)
    if fp_bits_len > 0:
        xor_fp_bits = wm_fp.extract(
            pirate_xor,
            n_bits=fp_bits_len,
            content_id=f"{args.content_id}{FP_SUFFIX}",
        )
    xor_scores = accuse_symmetric(xor_fp_bits, codebook, p)
    xor_ranking = rank_suspects(users, xor_scores)

    pirate_xor_paths = save_main_asset(
        out, "pirate_xor", pirate_xor, source_meta, emit_psb=args.emit_psb
    )
    if len(xor_fp_bits) > 0:
        save_rgb(out / "pirate_xor_bits.png", wm_fp.bits_to_noise_image(xor_fp_bits))


    plot_scores(
        out / "scores.png",
        {
            "Average collusion": avg_ranking,
            "XOR-like collusion": xor_ranking,
        },
        colluders=set(colluder_ids),
    )

    diff = cv2.absdiff(cover, watermarked[users[0].user_id])
    save_rgb(out / "diff_first_user.png", diff)

    report = {
        "config": {
            "image": args.image if args.image is not None else "synthetic",
            "num_users": args.num_users,
            "code_length": args.code_length,
            "display_bits_len": DISPLAY_BITS_LEN,
            "fingerprint_bits_len": args.code_length,
            "colluders": colluder_ids,
            "content_id": args.content_id,
            "master_key": args.master_key,
            "seed": args.seed,
            "emit_psb": args.emit_psb,
            "psb_max_side": args.psb_max_side,
            "input_is_psb": source_meta.input_is_psb,
            "input_depth": source_meta.depth,
            "delta_mode": args.delta_mode,
            "text_delta": max(14.0, float(args.delta) * 2.0),
            "text_repeats": max(8, int(args.repeats) * 4),
            "fp_delta": float(args.delta),
            "fp_repeats": int(args.repeats),
        },
        "clean_ber": clean_extract_ber,
        "psnr": psnr_map,
        "user_results": user_results,
        "average_attack": {
            "decoded_text": avg_text,
            "decoded_text_ok": avg_text_ok,
            "top5": avg_ranking[:5],
            "top3_hit_count": sum(uid in colluder_ids for uid, _ in avg_ranking[:3]),
        },
        "xor_attack": {
            "decoded_text": xor_text,
            "decoded_text_ok": xor_text_ok,
            "top5": xor_ranking[:5],
            "top3_hit_count": sum(uid in colluder_ids for uid, _ in xor_ranking[:3]),
        },
        "outputs": {
            "cover": cover_paths,
            "users": user_outputs,
            "pirate_average": pirate_avg_paths,
            "pirate_xor": pirate_xor_paths,
        },
        "width": int(source_meta.width),
        "height": int(source_meta.height),
        "input_is_psb": bool(source_meta.input_is_psb),
        "input_depth": int(source_meta.depth) if source_meta.depth is not None else None,
    }

    (out / "report.json").write_text(
        json.dumps(report, indent=2, ensure_ascii=False),
        encoding="utf-8",
    )
    return report

def build_watermarkers(
    master_key: str,
    delta: float,
    repeats: int,
    delta_mode: str,
) -> tuple[BlindDwtDctQim, BlindDwtDctQim]:
    """
    返回两个水印器：
    1) 字符串通道：更强保护
    2) 指纹通道：正常保护
    """
    wm_text = BlindDwtDctQim(
        master_key=master_key,
        delta=max(14.0, float(delta) * 2.0),
        repeats=max(8, int(repeats) * 4),
        coeff_pos=(2, 2),
        bands=("LH", "HL"),
        delta_mode=delta_mode,
    )

    wm_fp = BlindDwtDctQim(
        master_key=master_key,
        delta=float(delta),
        repeats=int(repeats),
        coeff_pos=(3, 3),
        bands=("LH", "HL"),
        delta_mode=delta_mode,
    )
    return wm_text, wm_fp


def build_single_user_channels(
    user_id: str,
    user_key: str,
    content_id: str,
    master_key: str,
    code_length: int,
    seed: int,
) -> tuple[str, np.ndarray, np.ndarray]:
    """
    返回：
    1) display_text
    2) text_bits
    3) fingerprint_bits
    """
    display_text = build_display_text(user_id, content_id)
    text_bits = encode_display_text(display_text)

    fingerprint_bits = np.zeros(0, dtype=np.uint8)
    if int(code_length) > 0:
        user = UserRecord(user_id=user_id, user_key=user_key)
        p = sample_tardos_probabilities(
            length=int(code_length),
            cutoff=0.08,
            seed=seed,
        )
        fingerprint_bits = build_codebook([user], p, master_key=master_key)[0]

    return display_text, text_bits, fingerprint_bits

def fit_preview_size(width: int, height: int, max_side: int = 2400) -> tuple[int, int]:
    width = int(width)
    height = int(height)
    if width <= 0 or height <= 0:
        raise ValueError("width/height 必须为正数")

    scale = min(max_side / max(width, height), 1.0)
    out_w = max(64, int(round(width * scale)))
    out_h = max(64, int(round(height * scale)))
    return out_w, out_h

def load_large_tif_preview_bgr(
    tif_path: str | Path,
    out_png_path: str | Path,
    max_side: int = 2400,
) -> np.ndarray:
    save_tiff_preview_png(tif_path, out_png_path, max_side=max_side)
    preview_bgr = cv2.imread(str(out_png_path), cv2.IMREAD_COLOR)
    if preview_bgr is None:
        raise RuntimeError(f"无法读取大图预览 PNG: {out_png_path}")
    return preview_bgr

def scale_preview_quad_to_fullres(
    quad_preview: np.ndarray,
    preview_w: int,
    preview_h: int,
    full_w: int,
    full_h: int,
) -> np.ndarray:
    quad_preview = np.asarray(quad_preview, dtype=np.float32).reshape(4, 2)
    sx = float(full_w) / float(preview_w)
    sy = float(full_h) / float(preview_h)

    quad_full = quad_preview.copy()
    quad_full[:, 0] *= sx
    quad_full[:, 1] *= sy
    return quad_full.astype(np.float32)

def build_base_rectify_matrix(
    full_w: int,
    full_h: int,
    expected_w: int,
    expected_h: int,
    quad_full: np.ndarray | None,
) -> np.ndarray:
    dst_rect = np.array(
        [
            [0, 0],
            [expected_w - 1, 0],
            [expected_w - 1, expected_h - 1],
            [0, expected_h - 1],
        ],
        dtype=np.float32,
    )

    if quad_full is not None:
        return cv2.getPerspectiveTransform(
            np.asarray(quad_full, dtype=np.float32),
            dst_rect,
        )

    sx = (expected_w - 1) / max(full_w - 1, 1)
    sy = (expected_h - 1) / max(full_h - 1, 1)
    return np.array(
        [
            [sx, 0.0, 0.0],
            [0.0, sy, 0.0],
            [0.0, 0.0, 1.0],
        ],
        dtype=np.float64,
    )

def build_rotate_scale_matrix(
    out_w: int,
    out_h: int,
    angle_deg: float,
    scale: float,
) -> np.ndarray:
    center = (out_w / 2.0, out_h / 2.0)
    M2 = cv2.getRotationMatrix2D(center, angle_deg, scale)
    M3 = np.vstack([M2, np.array([0.0, 0.0, 1.0], dtype=np.float64)])
    return M3

def create_empty_large_tiff(
    out_tif_path: str | Path,
    width: int,
    height: int,
    dtype: np.dtype,
    source_meta: SourceMeta,
) -> Path:
    out_tif_path = Path(out_tif_path)
    out_tif_path.parent.mkdir(parents=True, exist_ok=True)

    with tifffile.TiffWriter(str(out_tif_path), bigtiff=True) as tw:
        write_kwargs = {
            "data": None,
            "shape": (height, width, 3),
            "dtype": dtype,
            "photometric": "rgb",
            "compression": None,
            "metadata": None,
        }
        if source_meta.dpi_x and source_meta.dpi_y:
            write_kwargs["resolution"] = (source_meta.dpi_x, source_meta.dpi_y)
            write_kwargs["resolutionunit"] = "INCH"
        if source_meta.icc_profile:
            write_kwargs["iccprofile"] = source_meta.icc_profile
        tw.write(**write_kwargs)

    return out_tif_path

def warp_large_tiff_by_matrix(
    src_tif_path: str | Path,
    dst_tif_path: str | Path,
    out_w: int,
    out_h: int,
    M: np.ndarray,
    source_meta: SourceMeta,
    tile_size: int = 2048,
    border_mode: int = cv2.BORDER_REPLICATE,
) -> Path:
    src_tif_path = Path(src_tif_path)
    dst_tif_path = Path(dst_tif_path)

    src_mm = None
    dst_mm = None
    try:
        src_mm = memmap_tiff_rgb(src_tif_path, mode="r")
        src_h, src_w = int(src_mm.shape[0]), int(src_mm.shape[1])
        dtype = src_mm.dtype

        create_empty_large_tiff(dst_tif_path, out_w, out_h, dtype, source_meta)
        dst_mm = memmap_tiff_rgb(dst_tif_path, mode="r+")

        M = np.asarray(M, dtype=np.float64)
        M_inv = np.linalg.inv(M)

        for y0, y1, x0, x1 in iter_tiles(out_h, out_w, tile_size):
            tile_h = y1 - y0
            tile_w = x1 - x0

            dst_corners = np.array(
                [[[x0, y0], [x1 - 1, y0], [x1 - 1, y1 - 1], [x0, y1 - 1]]],
                dtype=np.float32,
            )
            src_corners = cv2.perspectiveTransform(dst_corners, M_inv.astype(np.float32))[0]

            sx0 = max(0, int(np.floor(np.min(src_corners[:, 0]))) - 8)
            sy0 = max(0, int(np.floor(np.min(src_corners[:, 1]))) - 8)
            sx1 = min(src_w, int(np.ceil(np.max(src_corners[:, 0]))) + 8)
            sy1 = min(src_h, int(np.ceil(np.max(src_corners[:, 1]))) + 8)

            if sx1 <= sx0 or sy1 <= sy0:
                continue

            patch = np.asarray(src_mm[sy0:sy1, sx0:sx1, :])

            T_src = np.array(
                [[1.0, 0.0, float(sx0)],
                 [0.0, 1.0, float(sy0)],
                 [0.0, 0.0, 1.0]],
                dtype=np.float64,
            )
            T_dst = np.array(
                [[1.0, 0.0, -float(x0)],
                 [0.0, 1.0, -float(y0)],
                 [0.0, 0.0, 1.0]],
                dtype=np.float64,
            )

            A = T_dst @ M @ T_src

            warped_tile = cv2.warpPerspective(
                patch,
                A.astype(np.float32),
                (tile_w, tile_h),
                flags=cv2.INTER_CUBIC,
                borderMode=border_mode,
            )

            dst_mm[y0:y1, x0:x1, :] = warped_tile

        return dst_tif_path
    finally:
        release_memmap(src_mm)
        release_memmap(dst_mm)

def save_large_asset_from_tif(
    out_dir: Path,
    stem: str,
    tif_path: str | Path,
    source_meta: SourceMeta,
    emit_psb: bool,
    preview_max_side: int = 2048,
) -> dict:
    tif_path = Path(tif_path)
    out_dir.mkdir(parents=True, exist_ok=True)

    final_tif = out_dir / f"{stem}.tif"
    if tif_path.resolve() != final_tif.resolve():
        clone_raster_file(tif_path, final_tif)

    png_path = out_dir / f"{stem}.png"
    save_tiff_preview_png(final_tif, png_path, max_side=preview_max_side)

    psb_path = None
    psb_error = None
    if emit_psb:
        candidate = out_dir / f"{stem}.psb"
        try:
            save_flattened_tiff_as_psb_via_photoshop(
                final_tif,
                candidate,
                dpi_x=source_meta.dpi_x,
                dpi_y=source_meta.dpi_y,
            )
            psb_path = candidate
        except Exception as exc:
            psb_error = str(exc)

    return {
        "png": str(png_path),
        "tif": str(final_tif),
        "psb": str(psb_path) if psb_path else None,
        "psb_error": psb_error,
    }

def iter_tiles(height: int, width: int, tile_size: int):
    for y0 in range(0, height, tile_size):
        y1 = min(y0 + tile_size, height)
        for x0 in range(0, width, tile_size):
            x1 = min(x0 + tile_size, width)
            yield y0, y1, x0, x1

def compute_tiled_mse(
    ref_tif_path: str | Path,
    test_tif_path: str | Path,
    tile_size: int = 2048,
) -> float:
    ref_mm = None
    test_mm = None
    try:
        ref_mm = memmap_tiff_rgb(ref_tif_path, mode="r")
        test_mm = memmap_tiff_rgb(test_tif_path, mode="r")

        if ref_mm.shape != test_mm.shape:
            raise ValueError(f"shape 不一致: {ref_mm.shape} vs {test_mm.shape}")
        if ref_mm.dtype != test_mm.dtype:
            raise ValueError(f"dtype 不一致: {ref_mm.dtype} vs {test_mm.dtype}")

        h, w = int(ref_mm.shape[0]), int(ref_mm.shape[1])
        sse = 0.0
        count = 0

        for y0, y1, x0, x1 in iter_tiles(h, w, tile_size):
            a = np.asarray(ref_mm[y0:y1, x0:x1, :], dtype=np.float32)
            b = np.asarray(test_mm[y0:y1, x0:x1, :], dtype=np.float32)
            diff = a - b
            sse += float(np.sum(diff * diff))
            count += int(diff.size)

        return 0.0 if count == 0 else sse / count
    finally:
        release_memmap(ref_mm)
        release_memmap(test_mm)

def compute_tiled_psnr(
    ref_tif_path: str | Path,
    test_tif_path: str | Path,
    tile_size: int = 2048,
) -> float:
    ref_mm = None
    try:
        ref_mm = memmap_tiff_rgb(ref_tif_path, mode="r")
        mse = compute_tiled_mse(ref_tif_path, test_tif_path, tile_size=tile_size)
        peak = dtype_peak(ref_mm.dtype)
        if mse <= 1e-12:
            return 99.0
        return float(10.0 * np.log10((peak ** 2) / mse))
    finally:
        release_memmap(ref_mm)

def average_collusion_tiled(
    colluder_tif_paths: list[str | Path],
    out_tif_path: str | Path,
    tile_size: int = 2048,
) -> Path:
    if not colluder_tif_paths:
        raise ValueError("colluder_tif_paths 不能为空")

    colluder_tif_paths = [Path(p) for p in colluder_tif_paths]
    out_tif_path = Path(out_tif_path)

    clone_raster_file(colluder_tif_paths[0], out_tif_path)

    src_mms = []
    out_mm = None
    try:
        src_mms = [memmap_tiff_rgb(p, mode="r") for p in colluder_tif_paths]
        out_mm = memmap_tiff_rgb(out_tif_path, mode="r+")

        dtype = src_mms[0].dtype
        peak = dtype_peak(dtype)
        h, w = int(out_mm.shape[0]), int(out_mm.shape[1])

        for mm in src_mms[1:]:
            if mm.shape != out_mm.shape:
                raise ValueError("所有合谋图 shape 必须一致")
            if mm.dtype != dtype:
                raise ValueError("所有合谋图 dtype 必须一致")

        for y0, y1, x0, x1 in iter_tiles(h, w, tile_size):
            acc = None
            for mm in src_mms:
                tile = np.asarray(mm[y0:y1, x0:x1, :], dtype=np.float32)
                acc = tile if acc is None else (acc + tile)

            avg_tile = np.clip(
                np.rint(acc / len(src_mms)),
                0,
                peak,
            ).astype(dtype)

            out_mm[y0:y1, x0:x1, :] = avg_tile

        return out_tif_path
    finally:
        release_memmap(out_mm)
        for mm in src_mms:
            release_memmap(mm)

def extract_payload_from_tif(
    tif_path: str | Path,
    wm: BlindDwtDctQim,
    n_bits: int,
    content_id: str,
    tile_size: int = 2048,
) -> np.ndarray:
    mm = None
    try:
        mm = memmap_tiff_rgb(tif_path, mode="r")
        return wm.extract_tiled_rgb(
            mm,
            n_bits=n_bits,
            content_id=content_id,
            tile_size=tile_size,
        )
    finally:
        release_memmap(mm)

def xor_collusion_via_reembed_tiled(
    base_tif_path: str | Path,
    colluder_tif_paths: list[str | Path],
    out_tif_path: str | Path,
    wm: BlindDwtDctQim,
    n_bits: int,
    content_id: str,
    tile_size: int = 2048,
) -> tuple[Path, np.ndarray]:
    if not colluder_tif_paths:
        raise ValueError("colluder_tif_paths 不能为空")

    extracted_list = []
    for p in colluder_tif_paths:
        bits = extract_payload_from_tif(
            tif_path=p,
            wm=wm,
            n_bits=n_bits,
            content_id=content_id,
            tile_size=tile_size,
        )
        extracted_list.append(bits)

    fused = np.bitwise_xor.reduce(np.stack(extracted_list, axis=0), axis=0).astype(np.uint8)

    base_tif_path = Path(base_tif_path)
    out_tif_path = Path(out_tif_path)
    clone_raster_file(base_tif_path, out_tif_path)

    out_mm = None
    try:
        out_mm = memmap_tiff_rgb(out_tif_path, mode="r+")
        wm.embed_tiled_rgb_inplace(
            out_mm,
            fused,
            content_id=content_id,
            tile_size=tile_size,
        )
        return out_tif_path, fused
    finally:
        release_memmap(out_mm)

def run_large_demo(args) -> dict:
    out = Path(args.out_dir)
    out.mkdir(parents=True, exist_ok=True)

    #work_dir = out / "_large_work"
    work_dir = out / "_large_work" / uuid.uuid4().hex
    work_dir.mkdir(parents=True, exist_ok=True)

    input_raster_path, source_meta = prepare_large_raster_input(
        args.image,
        out_dir=work_dir,
        tile_size=getattr(args, "tile_size", 2048),
    )

    tile_size = int(getattr(args, "tile_size", 2048))
    preview_max_side = 2048

    cover_paths = save_large_asset_from_tif(
        out_dir=out,
        stem="cover",
        tif_path=input_raster_path,
        source_meta=source_meta,
        emit_psb=args.emit_psb,
        preview_max_side=preview_max_side,
    )

    users = build_users(args.num_users)
    user_ids = {u.user_id for u in users}
    colluder_ids = parse_colluders(args.colluders, user_ids)

    p = sample_tardos_probabilities(
        length=args.code_length,
        cutoff=0.08,
        seed=args.seed,
    )
    codebook = build_codebook(users, p, master_key=args.master_key)

    wm_text, wm_fp = build_watermarkers(
        master_key=args.master_key,
        delta=args.delta,
        repeats=args.repeats,
        delta_mode=args.delta_mode,
    )

    text_bits_len = DISPLAY_BITS_LEN
    fp_bits_len = args.code_length

    clean_extract_ber = {}
    psnr_map = {}
    user_outputs = {}
    user_results = {}
    user_tif_paths = {}

    for idx, user in enumerate(users):
        display_text, text_bits, fingerprint_bits = build_single_user_channels(
            user_id=user.user_id,
            user_key=user.user_key,
            content_id=args.content_id,
            master_key=args.master_key,
            code_length=args.code_length,
            seed=args.seed,
        )

        user_tif = out / f"{user.user_id}_watermarked.tif"
        clone_raster_file(input_raster_path, user_tif)

        rgb_mm = None
        try:
            rgb_mm = memmap_tiff_rgb(user_tif, mode="r+")
            wm_text.embed_tiled_rgb_inplace(
                rgb_mm,
                text_bits,
                content_id=f"{args.content_id}{TEXT_SUFFIX}",
                tile_size=tile_size,
            )

            if len(fingerprint_bits) > 0:
                wm_fp.embed_tiled_rgb_inplace(
                    rgb_mm,
                    fingerprint_bits,
                    content_id=f"{args.content_id}{FP_SUFFIX}",
                    tile_size=tile_size,
                )
        finally:
            release_memmap(rgb_mm)

        read_mm = None
        try:
            read_mm = memmap_tiff_rgb(user_tif, mode="r")

            extracted_text_bits = wm_text.extract_tiled_rgb(
                read_mm,
                n_bits=len(text_bits),
                content_id=f"{args.content_id}{TEXT_SUFFIX}",
                tile_size=tile_size,
            )
            extracted_text, text_ok = decode_display_text(extracted_text_bits)

            extracted_fingerprint_bits = np.zeros(0, dtype=np.uint8)
            fingerprint_ber = None
            if len(fingerprint_bits) > 0:
                extracted_fingerprint_bits = wm_fp.extract_tiled_rgb(
                    read_mm,
                    n_bits=len(fingerprint_bits),
                    content_id=f"{args.content_id}{FP_SUFFIX}",
                    tile_size=tile_size,
                )
                fingerprint_ber = wm_fp.bit_error_rate(
                    fingerprint_bits,
                    extracted_fingerprint_bits,
                )
        finally:
            release_memmap(read_mm)

        clean_extract_ber[user.user_id] = fingerprint_ber if fingerprint_ber is not None else 0.0

        psnr_map[user.user_id] = compute_tiled_psnr(
            ref_tif_path=input_raster_path,
            test_tif_path=user_tif,
            tile_size=tile_size,
        )

        user_outputs[user.user_id] = save_large_asset_from_tif(
            out_dir=out,
            stem=f"{user.user_id}_watermarked",
            tif_path=user_tif,
            source_meta=source_meta,
            emit_psb=args.emit_psb,
            preview_max_side=preview_max_side,
        )

        user_results[user.user_id] = {
            "embedded_text": display_text,
            "extracted_text": extracted_text,
            "text_ok": text_ok,
            "fingerprint_ber": fingerprint_ber,
        }

        if len(extracted_fingerprint_bits) > 0:
            save_rgb(
                out / f"{user.user_id}_extracted_bits.png",
                wm_fp.bits_to_noise_image(extracted_fingerprint_bits),
            )

        user_tif_paths[user.user_id] = user_tif

    colluder_tif_paths = [user_tif_paths[cid] for cid in colluder_ids]


    pirate_average_tif = out / "pirate_average.tif"
    average_collusion_tiled(
        colluder_tif_paths=colluder_tif_paths,
        out_tif_path=pirate_average_tif,
        tile_size=tile_size,
    )

    avg_text_bits = extract_payload_from_tif(
        tif_path=pirate_average_tif,
        wm=wm_text,
        n_bits=text_bits_len,
        content_id=f"{args.content_id}{TEXT_SUFFIX}",
        tile_size=tile_size,
    )
    avg_text, avg_text_ok = decode_display_text(avg_text_bits)

    avg_fp_bits = np.zeros(fp_bits_len, dtype=np.uint8)
    if fp_bits_len > 0:
        avg_fp_bits = extract_payload_from_tif(
            tif_path=pirate_average_tif,
            wm=wm_fp,
            n_bits=fp_bits_len,
            content_id=f"{args.content_id}{FP_SUFFIX}",
            tile_size=tile_size,
        )

    avg_scores = accuse_symmetric(avg_fp_bits, codebook, p)
    avg_ranking = rank_suspects(users, avg_scores)

    pirate_avg_paths = save_large_asset_from_tif(
        out_dir=out,
        stem="pirate_average",
        tif_path=pirate_average_tif,
        source_meta=source_meta,
        emit_psb=args.emit_psb,
        preview_max_side=preview_max_side,
    )
    if len(avg_fp_bits) > 0:
        save_rgb(out / "pirate_average_bits.png", wm_fp.bits_to_noise_image(avg_fp_bits))


    pirate_xor_tif = out / "pirate_xor.tif"
    pirate_xor_tif, xor_bits_fused = xor_collusion_via_reembed_tiled(
        base_tif_path=pirate_average_tif,
        colluder_tif_paths=colluder_tif_paths,
        out_tif_path=pirate_xor_tif,
        wm=wm_fp,
        n_bits=fp_bits_len,
        content_id=f"{args.content_id}{FP_SUFFIX}",
        tile_size=tile_size,
    )

    xor_text_bits = extract_payload_from_tif(
        tif_path=pirate_xor_tif,
        wm=wm_text,
        n_bits=text_bits_len,
        content_id=f"{args.content_id}{TEXT_SUFFIX}",
        tile_size=tile_size,
    )
    xor_text, xor_text_ok = decode_display_text(xor_text_bits)

    xor_fp_bits = np.zeros(fp_bits_len, dtype=np.uint8)
    if fp_bits_len > 0:
        xor_fp_bits = extract_payload_from_tif(
            tif_path=pirate_xor_tif,
            wm=wm_fp,
            n_bits=fp_bits_len,
            content_id=f"{args.content_id}{FP_SUFFIX}",
            tile_size=tile_size,
        )

    xor_scores = accuse_symmetric(xor_fp_bits, codebook, p)
    xor_ranking = rank_suspects(users, xor_scores)

    pirate_xor_paths = save_large_asset_from_tif(
        out_dir=out,
        stem="pirate_xor",
        tif_path=pirate_xor_tif,
        source_meta=source_meta,
        emit_psb=args.emit_psb,
        preview_max_side=preview_max_side,
    )
    if len(xor_fp_bits) > 0:
        save_rgb(out / "pirate_xor_bits.png", wm_fp.bits_to_noise_image(xor_fp_bits))

    plot_scores(
        out / "scores.png",
        {
            "Average collusion": avg_ranking,
            "XOR-like collusion": xor_ranking,
        },
        colluders=set(colluder_ids),
    )

    report = {
        "mode": "large_multi_user_tiled",
        "config": {
            "image": args.image if args.image is not None else "synthetic",
            "num_users": args.num_users,
            "code_length": args.code_length,
            "display_bits_len": DISPLAY_BITS_LEN,
            "colluders": colluder_ids,
            "fingerprint_bits_len": args.code_length,
            "text_delta": max(14.0, float(args.delta) * 2.0),
            "text_repeats": max(8, int(args.repeats) * 4),
            "fp_delta": float(args.delta),
            "fp_repeats": int(args.repeats),
            "content_id": args.content_id,
            "master_key": args.master_key,
            "seed": args.seed,
            "emit_psb": args.emit_psb,
            "psb_max_side": args.psb_max_side,
            "input_is_psb": source_meta.input_is_psb,
            "input_depth": source_meta.depth,
            "delta_mode": args.delta_mode,
            "tile_size": tile_size,
        },
        "input_raster": str(input_raster_path),
        "clean_ber": clean_extract_ber,
        "psnr": psnr_map,
        "user_results": user_results,
        "average_attack": {
            "decoded_text": avg_text,
            "decoded_text_ok": avg_text_ok,
            "top5": avg_ranking[:5],
            "top3_hit_count": sum(uid in colluder_ids for uid, _ in avg_ranking[:3]),
        },
        "xor_attack": {
            "decoded_text": xor_text,
            "decoded_text_ok": xor_text_ok,
            "top5": xor_ranking[:5],
            "top3_hit_count": sum(uid in colluder_ids for uid, _ in xor_ranking[:3]),
        },
        "outputs": {
            "cover": cover_paths,
            "users": user_outputs,
            "pirate_average": pirate_avg_paths,
            "pirate_xor": pirate_xor_paths,
        },
        "tile_size": tile_size,
        "width": int(source_meta.width),
        "height": int(source_meta.height),
        "input_is_psb": bool(source_meta.input_is_psb),
        "input_depth": int(source_meta.depth) if source_meta.depth is not None else None,
    }

    (out / "report.json").write_text(
        json.dumps(report, indent=2, ensure_ascii=False),
        encoding="utf-8",
    )
    return report

def run_large_single_file(
    image_path: str,
    out_dir: str,
    user_id: str,
    user_key: str,
    content_id: str = "asset-demo-001",
    master_key: str = "server-master-key-2026",
    code_length: int = 512,
    seed: int = 2026,
    delta: float = 8.0,
    repeats: int = 2,
    delta_mode: str = "normalized_8bit",
    tile_size: int = 2048,
    emit_psb: bool = False,
    preview_max_side: int = 2048,
    expected_width: int | None = None,
    expected_height: int | None = None,
) -> dict:
    """针对超大 PSB/TIFF 的单文件分块水印入口。

    典型链路：PSB -> Photoshop 扁平导出 TIFF -> memmap 分块嵌入 -> TIFF/PNG -> (可选) PSB。
    这条链路不再要求把整张图一次性读进 numpy 内存。
    """
    out = Path(out_dir)
    out.mkdir(parents=True, exist_ok=True)
    #work_dir = out / "_large_work"
    work_dir = out / "_large_work" / uuid.uuid4().hex
    work_dir.mkdir(parents=True, exist_ok=True)

    embedded_text, text_bits, fingerprint_bits = build_single_user_channels(
        user_id=user_id,
        user_key=user_key,
        content_id=content_id,
        master_key=master_key,
        code_length=code_length,
        seed=seed,
    )

    input_raster_path, source_meta = prepare_large_raster_input(image_path, out_dir=work_dir, tile_size=tile_size)

    wm_text, wm_fp = build_watermarkers(
        master_key=master_key,
        delta=delta,
        repeats=repeats,
        delta_mode=delta_mode,
    )

    stem = Path(image_path).stem
    output_tif = out / f"{stem}_watermarked_large.tif"
    clone_raster_file(input_raster_path, output_tif)

    rgb_mm = None
    try:
        rgb_mm = memmap_tiff_rgb(output_tif, mode="r+")
        text_jobs = wm_text.embed_tiled_rgb_inplace(
            rgb_mm,
            text_bits,
            content_id=f"{content_id}{TEXT_SUFFIX}",
            tile_size=tile_size,
        )

        fp_jobs = []
        if len(fingerprint_bits) > 0:
            fp_jobs = wm_fp.embed_tiled_rgb_inplace(
                rgb_mm,
                fingerprint_bits,
                content_id=f"{content_id}{FP_SUFFIX}",
                tile_size=tile_size,
            )
    finally:
        release_memmap(rgb_mm)

    preview_png = out / f"{stem}_watermarked_large.png"
    save_tiff_preview_png(output_tif, preview_png, max_side=preview_max_side)

    output_psb = None
    psb_error = None
    if emit_psb:
        output_psb = out / f"{stem}_watermarked_large.psb"
        try:
            save_flattened_tiff_as_psb_via_photoshop(
                output_tif,
                output_psb,
                dpi_x=source_meta.dpi_x,
                dpi_y=source_meta.dpi_y,
            )
        except Exception as exc:
            psb_error = str(exc)
            output_psb = None

    read_mm = None
    try:
        read_mm = memmap_tiff_rgb(output_tif, mode="r")

        extracted_text_bits = wm_text.extract_tiled_rgb(
            read_mm,
            n_bits=len(text_bits),
            content_id=f"{content_id}{TEXT_SUFFIX}",
            tile_size=tile_size,
        )
        extracted_text, text_ok = decode_display_text(extracted_text_bits)

        extracted_fingerprint = np.zeros(0, dtype=np.uint8)
        fingerprint_ber = None
        if len(fingerprint_bits) > 0:
            extracted_fingerprint = wm_fp.extract_tiled_rgb(
                read_mm,
                n_bits=len(fingerprint_bits),
                content_id=f"{content_id}{FP_SUFFIX}",
                tile_size=tile_size,
            )
            fingerprint_ber = wm_fp.bit_error_rate(
                fingerprint_bits,
                extracted_fingerprint,
            )
    finally:
        release_memmap(read_mm)


    report = {
        "mode": "large_single_file_tiled",
        "image": image_path,
        "input_raster": str(input_raster_path),
        "output_tif": str(output_tif),
        "output_png": str(preview_png),
        "output_psb": str(output_psb) if output_psb else None,
        "output_psb_error": psb_error,
        "user_id": user_id,
        "content_id": content_id,
        "embedded_text": embedded_text,
        "extracted_text": extracted_text,
        "text_ok": bool(text_ok),
        "fingerprint_ber": fingerprint_ber,
        "tile_size": int(tile_size),
        "text_tile_count_used": len(text_jobs),
        "fp_tile_count_used": len(fp_jobs),
        "display_bits_len": int(len(text_bits)),
        "fingerprint_bits_len": int(len(fingerprint_bits)),
        "text_delta": float(max(14.0, delta * 2.0)),
        "text_repeats": int(max(8, repeats * 4)),
        "fp_delta": float(delta),
        "fp_repeats": int(repeats),
        "delta_mode": delta_mode,
        "input_is_psb": bool(source_meta.input_is_psb),
        "input_depth": int(source_meta.depth) if source_meta.depth is not None else None,
        "width": int(source_meta.width),
        "height": int(source_meta.height),
    }

    (out / "large_report.json").write_text(
        json.dumps(report, indent=2, ensure_ascii=False),
        encoding="utf-8",
    )
    return report


def extract_large_single_file(
    image_path: str,
    out_dir: str,
    content_id: str = "asset-demo-001",
    master_key: str = "server-master-key-2026",
    code_length: int = 512,
    delta: float = 8.0,
    repeats: int = 2,
    delta_mode: str = "normalized_8bit",
    tile_size: int = 2048,
    expected_width: int | None = None,
    expected_height: int | None = None,
) -> dict:
    out = Path(out_dir)
    out.mkdir(parents=True, exist_ok=True)
    work_dir = out / "_large_extract_work" / uuid.uuid4().hex
    work_dir.mkdir(parents=True, exist_ok=True)

    input_raster_path, source_meta = prepare_large_raster_input(image_path, out_dir=work_dir, tile_size=tile_size)

    preview_png = work_dir / "scan_preview.png"
    preview_bgr = load_large_tif_preview_bgr(
        input_raster_path,
        preview_png,
        max_side=2400,
    )

    preview_h, preview_w = preview_bgr.shape[:2]
    full_w = int(source_meta.width)
    full_h = int(source_meta.height)

    quad_preview = None
    quad_full = None
    base_rectify_M = None

    if expected_width and expected_height:
        quad_preview = detect_page_quad(preview_bgr)
        if quad_preview is not None:
            quad_full = scale_preview_quad_to_fullres(
                quad_preview,
                preview_w=preview_w,
                preview_h=preview_h,
                full_w=full_w,
                full_h=full_h,
            )

        base_rectify_M = build_base_rectify_matrix(
            full_w=full_w,
            full_h=full_h,
            expected_w=int(expected_width),
            expected_h=int(expected_height),
            quad_full=quad_full,
        )

    wm_text, _ = build_watermarkers(
        master_key=master_key,
        delta=delta,
        repeats=repeats,
        delta_mode=delta_mode,
    )

    best_candidate = None
    extract_raster_path = input_raster_path
    candidate_index = 0

    if expected_width and expected_height and base_rectify_M is not None:
        angle_candidates = [-1.0, -0.5, 0.0, 0.5, 1.0]
        scale_candidates = [0.99, 1.0, 1.01]

        for angle in angle_candidates:
            for scale in scale_candidates:
                rot_M = build_rotate_scale_matrix(
                    out_w=int(expected_width),
                    out_h=int(expected_height),
                    angle_deg=float(angle),
                    scale=float(scale),
                )
                total_M = rot_M @ base_rectify_M

                candidate_tif = work_dir / f"scan_candidate_{candidate_index:02d}.tif"
                warp_large_tiff_by_matrix(
                    src_tif_path=input_raster_path,
                    dst_tif_path=candidate_tif,
                    out_w=int(expected_width),
                    out_h=int(expected_height),
                    M=total_M,
                    source_meta=source_meta,
                    tile_size=tile_size,
                )

                cand = score_large_text_candidate(
                    tif_path=candidate_tif,
                    wm_text=wm_text,
                    content_id=content_id,
                    tile_size=tile_size,
                )
                cand["candidate_tif"] = str(candidate_tif)
                cand["angle"] = float(angle)
                cand["scale"] = float(scale)
                cand["used_quad"] = quad_full is not None

                if best_candidate is None or cand["score"] > best_candidate["score"]:
                    best_candidate = cand

                candidate_index += 1

        if best_candidate is not None:
            extract_raster_path = Path(best_candidate["candidate_tif"])

    if best_candidate is not None:
        extracted_text_bits = best_candidate["bits"]
        extracted_text = best_candidate["extracted_text"]
        text_ok = bool(best_candidate["text_ok"])
        soft = best_candidate["soft"]
    else:
        read_mm = None
        try:
            read_mm = memmap_tiff_rgb(extract_raster_path, mode="r")
            extracted_text_bits, soft = wm_text.extract_tiled_rgb_with_confidence(
                read_mm,
                n_bits=DISPLAY_BITS_LEN,
                content_id=f"{content_id}{TEXT_SUFFIX}",
                tile_size=tile_size,
            )
            extracted_text, text_ok = decode_display_text(extracted_text_bits)
        finally:
            release_memmap(read_mm)

    report = {
        "mode": "large_single_file_extract_tiled",
        "image": image_path,
        "input_raster": str(input_raster_path),
        "content_id": content_id,
        "display_bits_len": int(DISPLAY_BITS_LEN),
        "extracted_text": extracted_text,
        "text_ok": bool(text_ok),
        "tile_size": int(tile_size),
        "text_delta": float(max(14.0, delta * 2.0)),
        "text_repeats": int(max(8, repeats * 4)),
        "fp_delta": float(delta),
        "fp_repeats": int(repeats),
        "delta_mode": delta_mode,
        "input_is_psb": bool(source_meta.input_is_psb),
        "input_depth": int(source_meta.depth) if source_meta.depth is not None else None,
        "width": int(source_meta.width),
        "height": int(source_meta.height),
        "scan_enhanced": bool(expected_width and expected_height),
        "scan_used_quad": bool(best_candidate["used_quad"]) if best_candidate else False,
        "scan_best_angle": best_candidate["angle"] if best_candidate else 0.0,
        "scan_best_scale": best_candidate["scale"] if best_candidate else 1.0,
        "scan_candidate_count": int(candidate_index) if expected_width and expected_height else 0,
        "normalized_extract_raster": str(extract_raster_path),
    }
    (out / "large_extract_report.json").write_text(
        json.dumps(report, indent=2, ensure_ascii=False),
        encoding="utf-8",
    )
    return report

def score_large_text_candidate(
    tif_path: str | Path,
    wm_text,
    content_id: str,
    tile_size: int,
) -> dict:
    mm = None
    try:
        mm = memmap_tiff_rgb(tif_path, mode="r")
        bits, soft = wm_text.extract_tiled_rgb_with_confidence(
            mm,
            n_bits=DISPLAY_BITS_LEN,
            content_id=f"{content_id}{TEXT_SUFFIX}",
            tile_size=tile_size,
        )
        extracted_text, text_ok = decode_display_text(bits)
        score = (1e9 if text_ok else 0.0) + float(np.sum(np.abs(soft)))
        return {
            "bits": bits,
            "soft": soft,
            "extracted_text": extracted_text,
            "text_ok": bool(text_ok),
            "score": score,
        }
    finally:
        release_memmap(mm)

if __name__ == "__main__":
    args = parse_args()
    report = run_demo(args)
    print(json.dumps(report, indent=2, ensure_ascii=False))