import csv
import shutil

from config import ORIGINAL_DIR, WATERMARKED_DIR, ATTACKED_DIR, VERIFY_CHAIN_CSV, PLATFORM_NAME
from common_image import (
    embed_watermark, sha256_file, phash_file, phash_distance,
    extract_watermark, attack_jpeg, attack_scale, attack_crop, attack_screenshot_like
)
from common_chain import get_contract, get_account, register_record, get_all_records_from_chain, estimate_query_gas, query_record


PHASH_THRESHOLD = 15
REPRESENTATIVE_IMAGE = "zuixuanminzufeng.png"


def verify_one_image(img_path, chain_records):
    current_sha = sha256_file(img_path)
    current_phash = phash_file(img_path)

    for r in chain_records:
        if current_sha == r["content_id"]:
            return {
                "status": "EXACT_MATCH",
                "matched_file": r["file_name"],
                "matched_wm": r["watermark_id"],
                "distance": 0,
                "extracted_wm": ""
            }

    wm = extract_watermark(img_path, wm_len_bytes=8)
    if wm:
        for r in chain_records:
            if wm == r["watermark_id"]:
                dist = phash_distance(current_phash, r["phash_value"])
                if dist <= PHASH_THRESHOLD:
                    return {
                        "status": "SIMILAR_VARIANT",
                        "matched_file": r["file_name"],
                        "matched_wm": r["watermark_id"],
                        "distance": dist,
                        "extracted_wm": wm
                    }
                else:
                    return {
                        "status": "TAMPERED",
                        "matched_file": r["file_name"],
                        "matched_wm": r["watermark_id"],
                        "distance": dist,
                        "extracted_wm": wm
                    }

    best_record = None
    best_dist = None
    for r in chain_records:
        dist = phash_distance(current_phash, r["phash_value"])
        if best_dist is None or dist < best_dist:
            best_dist = dist
            best_record = r

    if best_record is not None and best_dist <= PHASH_THRESHOLD:
        return {
            "status": "SIMILAR_VARIANT",
            "matched_file": best_record["file_name"],
            "matched_wm": best_record["watermark_id"],
            "distance": best_dist,
            "extracted_wm": wm
        }

    return {
        "status": "NOT_FOUND",
        "matched_file": "",
        "matched_wm": "",
        "distance": best_dist if best_dist is not None else -1,
        "extracted_wm": wm
    }


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

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

    original_path = ORIGINAL_DIR / REPRESENTATIVE_IMAGE
    if not original_path.exists():
        print(f"未找到代表图片: {original_path}")
        return

    wm_id = "WM900001"
    wm_path = WATERMARKED_DIR / f"{original_path.stem}_paper_wm.png"
    WATERMARKED_DIR.mkdir(parents=True, exist_ok=True)

    embed_watermark(original_path, wm_path, wm_id)
    content_id = sha256_file(wm_path)
    phash_value = phash_file(wm_path)

    reg_result = register_record(
        contract, w3, account,
        content_id,
        wm_path.name,
        PLATFORM_NAME,
        wm_id,
        phash_value
    )

    print("[代表样本上链成功]")
    print("content_id =", content_id)
    print("tx_hash    =", reg_result["tx_hash"])
    print("gas_used   =", reg_result["gas_used"])
    print("block_no   =", reg_result["block_number"])

    query_gas = estimate_query_gas(contract, account.address, content_id)
    chain_record = query_record(contract, content_id)

    print("\n[查询结果]")
    print(chain_record)
    print("estimated query gas =", query_gas)

    case_dir = ATTACKED_DIR / f"{original_path.stem}_paper"
    case_dir.mkdir(parents=True, exist_ok=True)

    exact_path = case_dir / f"{original_path.stem}_exact.png"
    jpeg_path = case_dir / f"{original_path.stem}_jpeg70.jpg"
    scale_path = case_dir / f"{original_path.stem}_scale50.png"
    crop_path = case_dir / f"{original_path.stem}_crop80.png"
    shot_path = case_dir / f"{original_path.stem}_shot.png"

    shutil.copy(wm_path, exact_path)
    attack_jpeg(wm_path, jpeg_path, quality=70)
    attack_scale(wm_path, scale_path, scale=0.5)
    attack_crop(wm_path, crop_path, ratio=0.8)
    attack_screenshot_like(wm_path, shot_path)

    chain_records = get_all_records_from_chain(contract)

    test_images = {
        "exact": exact_path,
        "jpeg70": jpeg_path,
        "scale50": scale_path,
        "crop80": crop_path,
        "shot": shot_path
    }

    with VERIFY_CHAIN_CSV.open("w", newline="", encoding="utf-8") as f:
        writer = csv.writer(f)
        writer.writerow([
            "attack_type", "test_image", "status",
            "matched_file", "matched_wm", "extracted_wm", "phash_distance"
        ])

        for attack_type, img_path in test_images.items():
            result = verify_one_image(img_path, chain_records)

            print("-" * 60)
            print("attack_type =", attack_type)
            print("status      =", result["status"])
            print("matched     =", result["matched_file"])
            print("wm          =", result["extracted_wm"])
            print("distance    =", result["distance"])

            writer.writerow([
                attack_type,
                img_path.name,
                result["status"],
                result["matched_file"],
                result["matched_wm"],
                result["extracted_wm"],
                result["distance"]
            ])

    print(f"\n链路验证结果已写入: {VERIFY_CHAIN_CSV}")


if __name__ == "__main__":
    main()
