#!/usr/bin/env python3
"""
pdf_image_compare.py

Walk two PDFs, decode every image & soft mask, and print a comparison table
with per-object SHA-256s. Also writes CSVs for each PDF and a match report.

Usage:
  python3 pdf_image_compare.py left.pdf right.pdf

Dependencies:
  - pikepdf >= 7.0
  - Pillow (PIL) is NOT required for hashing; we hash the decoded image bytes
    directly from the PDF streams to avoid color-management differences.
"""
import sys
import csv
import hashlib
from dataclasses import dataclass, asdict
from typing import Dict, List, Optional, Tuple, Set

import pikepdf

@dataclass
class ImgRec:
    pdf_path: str
    page_index: Optional[int]  # zero-based, may be None if not resolved
    objnum: int
    gennum: int
    xobject_name: Optional[str]  # e.g., /Im1 if available
    width: int
    height: int
    bpc: Optional[int]
    colorspace: str
    filters: str
    has_smask: bool
    pixel_len: int
    pixel_sha256: str
    smask_len: Optional[int]
    smask_sha256: Optional[str]

def _name(obj) -> str:
    try:
        return str(obj) if obj is not None else ""
    except Exception:
        return ""

def _norm_filters(filt_obj) -> str:
    if filt_obj is None:
        return ""
    try:
        if isinstance(filt_obj, pikepdf.Name):
            return str(filt_obj)
        if isinstance(filt_obj, pikepdf.Array):
            return ",".join(str(x) for x in filt_obj)
        # Sometimes filters are indirect; stringify
        return str(filt_obj)
    except Exception:
        return str(filt_obj)

def sha256_hex(b: bytes) -> str:
    return hashlib.sha256(b).hexdigest()

def decode_stream_bytes(stream: pikepdf.Object) -> bytes:
    """
    Return the decoded (unfiltered) bytes of a PDF stream.
    This lets us hash the actual sample data independent of compression filter.
    """
    # pikepdf v7+: stream.read_bytes(decode_stream=True)
    try:
        return stream.read_bytes(decode_stream=True)
    except TypeError:
        # Fallback for older pikepdf
        return pikepdf.Stream(stream).read_bytes()

def walk_xobjects(pdf: pikepdf.Pdf, resources, seen: Set[Tuple[int,int]], pdf_path: str,
                  page_index: Optional[int], out: List[ImgRec]) -> None:
    if not resources:
        return
    xobj_dict = resources.get("/XObject", None)
    if not xobj_dict:
        return
    # Iterate over named XObjects in the Resources
    for name, xobj in xobj_dict.items():
        try:
            xobj = xobj.get_object()
        except Exception:
            continue
        if not isinstance(xobj, pikepdf.Stream):
            continue
        subtype = xobj.get("/Subtype", None)
        if subtype == pikepdf.Name("/Image"):
            # De-duplicate by object id
            key = (xobj.objgen[0], xobj.objgen[1])
            if key in seen:
                continue
            seen.add(key)
            width = int(xobj.get("/Width", 0))
            height = int(xobj.get("/Height", 0))
            bpc = int(xobj.get("/BitsPerComponent", 0)) if xobj.get("/BitsPerComponent", None) is not None else None
            colorspace = _name(xobj.get("/ColorSpace", ""))
            filters = _norm_filters(xobj.get("/Filter", None))

            # Decode pixel bytes (after filters)
            pixel_bytes = decode_stream_bytes(xobj)
            pixel_sha = sha256_hex(pixel_bytes) if pixel_bytes is not None else ""
            pixel_len = len(pixel_bytes) if pixel_bytes is not None else 0

            # Soft mask if present
            smask = xobj.get("/SMask", None)
            smask_len = None
            smask_sha = None
            if isinstance(smask, pikepdf.Object) or isinstance(smask, pikepdf.Stream):
                try:
                    smask = smask.get_object()
                except Exception:
                    pass
            if isinstance(smask, pikepdf.Stream):
                try:
                    smask_bytes = decode_stream_bytes(smask)
                    smask_sha = sha256_hex(smask_bytes)
                    smask_len = len(smask_bytes)
                except Exception:
                    smask_sha = None
                    smask_len = None

            rec = ImgRec(
                pdf_path=pdf_path,
                page_index=page_index,
                objnum=xobj.objgen[0],
                gennum=xobj.objgen[1],
                xobject_name=str(name),
                width=width,
                height=height,
                bpc=bpc,
                colorspace=colorspace,
                filters=filters,
                has_smask=smask_sha is not None,
                pixel_len=pixel_len,
                pixel_sha256=pixel_sha,
                smask_len=smask_len,
                smask_sha256=smask_sha,
            )
            out.append(rec)

        elif subtype == pikepdf.Name("/Form"):
            # Recurse into form XObject resources
            form_resources = xobj.get("/Resources", None)
            walk_xobjects(pdf, form_resources, seen, pdf_path, page_index, out)

def extract_images(pdf_path: str) -> List[ImgRec]:
    out: List[ImgRec] = []
    seen: Set[Tuple[int,int]] = set()
    with pikepdf.open(pdf_path) as pdf:
        # Walk pages
        for i, page in enumerate(pdf.pages):
            resources = page.get("/Resources", None)
            walk_xobjects(pdf, resources, seen, pdf_path, i, out)
    return out

def write_csv(rows: List[ImgRec], csv_path: str) -> None:
    fieldnames = list(asdict(rows[0]).keys()) if rows else [
        "pdf_path","page_index","objnum","gennum","xobject_name","width","height","bpc",
        "colorspace","filters","has_smask","pixel_len","pixel_sha256","smask_len","smask_sha256"
    ]
    with open(csv_path, "w", newline="", encoding="utf-8") as f:
        w = csv.DictWriter(f, fieldnames=fieldnames)
        w.writeheader()
        for r in rows:
            w.writerow(asdict(r))

def compare_sets(left: List[ImgRec], right: List[ImgRec]) -> Tuple[List[Dict], List[Dict]]:
    """
    Return (pixel_matches, smask_matches) as lists of dict rows describing matches.
    A match is defined as identical SHA-256 across PDFs.
    """
    L = {}
    for r in left:
        L.setdefault(r.pixel_sha256, []).append(r)
    R = {}
    for r in right:
        R.setdefault(r.pixel_sha256, []).append(r)
    pixel_matches = []
    for h in sorted(set(L.keys()) & set(R.keys())):
        for lrec in L[h]:
            for rrec in R[h]:
                pixel_matches.append({
                    "match_type": "PIXELS",
                    "left_pdf": lrec.pdf_path,
                    "left_obj": f"{lrec.objnum} {lrec.gennum}",
                    "left_page": lrec.page_index,
                    "left_dims": f"{lrec.width}x{lrec.height}",
                    "right_pdf": rrec.pdf_path,
                    "right_obj": f"{rrec.objnum} {rrec.gennum}",
                    "right_page": rrec.page_index,
                    "right_dims": f"{rrec.width}x{rrec.height}",
                    "sha256": h,
                })
    # Soft mask matches
    Ls = {}
    for r in left:
        if r.smask_sha256:
            Ls.setdefault(r.smask_sha256, []).append(r)
    Rs = {}
    for r in right:
        if r.smask_sha256:
            Rs.setdefault(r.smask_sha256, []).append(r)
    smask_matches = []
    for h in sorted(set(Ls.keys()) & set(Rs.keys())):
        for lrec in Ls[h]:
            for rrec in Rs[h]:
                smask_matches.append({
                    "match_type": "SMASK",
                    "left_pdf": lrec.pdf_path,
                    "left_obj": f"{lrec.objnum} {lrec.gennum}",
                    "left_page": lrec.page_index,
                    "left_dims": f"{lrec.width}x{lrec.height}",
                    "right_pdf": rrec.pdf_path,
                    "right_obj": f"{rrec.objnum} {rrec.gennum}",
                    "right_page": rrec.page_index,
                    "right_dims": f"{rrec.width}x{rrec.height}",
                    "sha256": h,
                })
    return pixel_matches, smask_matches

def main():
    if len(sys.argv) != 3:
        print("Usage: python3 pdf_image_compare.py left.pdf right.pdf")
        sys.exit(1)
    left_path, right_path = sys.argv[1], sys.argv[2]
    print(f"[+] Scanning images in: {left_path}")
    left = extract_images(left_path)
    print(f"[+] Found {len(left)} unique image objects in left PDF")
    print(f"[+] Scanning images in: {right_path}")
    right = extract_images(right_path)
    print(f"[+] Found {len(right)} unique image objects in right PDF")

    # Write per-PDF inventories
    left_csv = "left_images.csv"
    right_csv = "right_images.csv"
    write_csv(left, left_csv)
    write_csv(right, right_csv)
    print(f"[+] Wrote {left_csv} and {right_csv}")

    # Compare by decoded pixel bytes and soft-mask bytes
    pixel_matches, smask_matches = compare_sets(left, right)

    # Write match reports
    with open("matches_pixels.csv", "w", newline="", encoding="utf-8") as f:
        w = csv.DictWriter(f, fieldnames=[
            "match_type","left_pdf","left_obj","left_page","left_dims",
            "right_pdf","right_obj","right_page","right_dims","sha256"
        ])
        w.writeheader()
        for row in pixel_matches:
            w.writerow(row)
    with open("matches_smasks.csv", "w", newline="", encoding="utf-8") as f:
        w = csv.DictWriter(f, fieldnames=[
            "match_type","left_pdf","left_obj","left_page","left_dims",
            "right_pdf","right_obj","right_page","right_dims","sha256"
        ])
        w.writeheader()
        for row in smask_matches:
            w.writerow(row)

    # Console summary
    print(f"[=] Pixel-identical image matches: {len(pixel_matches)}")
    print(f"[=] Soft-mask (alpha) identical matches: {len(smask_matches)}")
    if pixel_matches:
        print("    Example match:")
        ex = pixel_matches[0]
        print("    ", ex)

if __name__ == "__main__":
    main()