from pathlib import Path
import shutil

from config import WATERMARKED_DIR, ATTACKED_DIR
from common_image import attack_jpeg, attack_scale, attack_crop, attack_screenshot_like


def main():
    ATTACKED_DIR.mkdir(parents=True, exist_ok=True)

    png_files = sorted(WATERMARKED_DIR.glob("*.png"))
    if not png_files:
        print("ex4_watermarked 中没有 PNG 文件，请先运行 ex4_batch_register_fullchain.py")
        return

    for img_path in png_files:
        stem = img_path.stem
        out_dir = ATTACKED_DIR / stem
        out_dir.mkdir(parents=True, exist_ok=True)

        # 复制一份 exact 样本，方便测试 EXACT_MATCH
        exact_path = out_dir / f"{stem}_exact.png"
        shutil.copy(img_path, exact_path)

        attack_jpeg(img_path, out_dir / f"{stem}_jpeg70.jpg", quality=70)
        attack_scale(img_path, out_dir / f"{stem}_scale50.png", scale=0.5)
        attack_crop(img_path, out_dir / f"{stem}_crop80.png", ratio=0.8)
        attack_screenshot_like(img_path, out_dir / f"{stem}_shot.png")

        print(f"[攻击生成完成] {img_path.name} -> {out_dir}")

    print("\n全部攻击样本已生成。")


if __name__ == "__main__":
    main()