import csv
from pathlib import Path

from config import ORIGINAL_DIR, WATERMARKED_DIR, PLATFORM_NAME, REGISTER_CSV
from common_image import embed_watermark, sha256_file, phash_file
from common_chain import get_contract, get_account, register_record


def main():
    WATERMARKED_DIR.mkdir(parents=True, exist_ok=True)
    REGISTER_CSV.parent.mkdir(parents=True, exist_ok=True)

    w3, contract = get_contract()
    account = get_account(w3)

    png_files = sorted(ORIGINAL_DIR.glob("*.png"))
    if not png_files:
        print("originals 目录中没有 PNG 文件。")
        return

    with REGISTER_CSV.open("w", newline="", encoding="utf-8") as f:
        writer = csv.writer(f)
        writer.writerow([
            "original_file", "watermarked_file", "watermark_id",
            "content_id", "phash_value", "tx_hash"
        ])

        for idx, img_path in enumerate(png_files, start=1):
            watermark_id = f"WM{idx:06d}"   # 固定 8 个字符
            wm_path = WATERMARKED_DIR / f"{img_path.stem}_wm.png"

            embed_watermark(img_path, wm_path, watermark_id)

            content_id = sha256_file(wm_path)
            phash_value = phash_file(wm_path)

            try:
                tx_hash = register_record(
                    contract, w3, account,
                    content_id,
                    wm_path.name,
                    PLATFORM_NAME,
                    watermark_id,
                    phash_value
                )
                print(f"[上链成功] {wm_path.name}  wm={watermark_id}  tx={tx_hash}")
            except Exception as e:
                tx_hash = f"ERROR: {e}"
                print(f"[上链失败] {wm_path.name}  wm={watermark_id}  err={e}")

            writer.writerow([
                img_path.name, wm_path.name, watermark_id,
                content_id, phash_value, tx_hash
            ])

    print(f"\n批量注册完成，结果已写入: {REGISTER_CSV}")


if __name__ == "__main__":
    main()