import csv

from config import ORIGINAL_DIR, WATERMARKED_DIR, PLATFORM_NAME, REGISTER_COST_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_COST_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_COST_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", "gas_used", "block_number"
        ])

        for idx, img_path in enumerate(png_files, start=1):
            watermark_id = f"WM{idx:06d}"
            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)

            result = register_record(
                contract, w3, account,
                content_id,
                wm_path.name,
                PLATFORM_NAME,
                watermark_id,
                phash_value
            )

            print(f"[上链成功] {wm_path.name}  tx={result['tx_hash']}  gas={result['gas_used']}")

            writer.writerow([
                img_path.name,
                wm_path.name,
                watermark_id,
                content_id,
                phash_value,
                result["tx_hash"],
                result["gas_used"],
                result["block_number"]
            ])

    print(f"\n上链成本记录已写入: {REGISTER_COST_CSV}")


if __name__ == "__main__":
    main()
