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

Finds *all* images in two PDFs — including INLINE images (BI ... ID ... EI)
and XObject images — decodes their bytes, and compares SHA-256 hashes.

It uses pdfminer.six to walk page layouts and obtain decoded image bytes,
which works for both inline and XObject images. It also tries to enrich
records with object numbers via pikepdf when possible.

Outputs:
  - left_images.csv, right_images.csv (per-PDF inventories)
  - matches_pixels.csv, matches_smasks.csv (identical decoded image / mask matches)

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

Install:
  pip install pdfminer.six pikepdf
"""
import sys
import csv
import hashlib
from dataclasses import dataclass, asdict
from typing import List, Dict, Tuple, Optional

# pdfminer imports
from pdfminer.high_level import extract_pages
from pdfminer.layout import LTImage, LAParams, LTContainer, LTAnno

# pikepdf (optional enrichment: object ids, filters, colorspace)
import pikepdf

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

@dataclass
class ImgRec:
    pdf_path: str
    page_index: int
    objnum: Optional[int]
    gennum: Optional[int]
    xobject_name: Optional[str]
    width: Optional[int]
    height: Optional[int]
    bpc: Optional[int]
    colorspace: Optional[str]
    filters: Optional[str]
    has_smask: Optional[bool]
    pixel_len: int
    pixel_sha256: str
    smask_len: Optional[int]
    smask_sha256: Optional[str]
    source: str  # 'inline' or 'xobject' (as seen by pdfminer)

def _flatten(layout):
    for obj in layout:
        if isinstance(obj, LTContainer):
            yield from _flatten(obj)
        else:
            yield obj

def _image_dims_from_stream(stream_dict) -> Tuple[Optional[int], Optional[int], Optional[int]]:
    try:
        w = stream_dict.get("Width", None)
        h = stream_dict.get("Height", None)
        bpc = stream_dict.get("BitsPerComponent", None)
        if isinstance(w, int) and isinstance(h, int):
            return w, h, bpc if isinstance(bpc, int) else None
    except Exception:
        pass
    return None, None, None

def inventory_with_pdfminer(pdf_path: str) -> List[ImgRec]:
    laparams = LAParams()
    out: List[ImgRec] = []
    # Use pikepdf to gather object metadata for enrichment
    try:
        ppdf = pikepdf.open(pdf_path)
    except Exception:
        ppdf = None

    for page_index, layout in enumerate(extract_pages(pdf_path, laparams=laparams)):
        for obj in _flatten(layout):
            if isinstance(obj, LTImage):
                # pdfminer LTImage gives decoded raw data via stream.get_data()
                try:
                    data = obj.stream.get_data()
                except Exception:
                    # As a fallback, use raw stream if decode fails
                    try:
                        data = obj.stream.get_rawdata()
                    except Exception:
                        continue

                pixel_sha = sha256_hex(data)
                pixel_len = len(data)

                # Soft mask: pdfminer may expose SMask as 'SMask' key in the dict
                smask_sha = None
                smask_len = None
                try:
                    smask = obj.stream.get("SMask")
                    if smask is not None and hasattr(smask, "get_data"):
                        sm_data = smask.get_data()
                        smask_sha = sha256_hex(sm_data)
                        smask_len = len(sm_data)
                except Exception:
                    pass

                objnum = None
                gennum = None
                xname = getattr(obj, "name", None)  # may look like 'Im0'
                width, height, bpc = _image_dims_from_stream(obj.stream)

                colorspace = None
                filters = None
                has_smask = smask_sha is not None

                # Enrich via pikepdf by searching matching image XObject on same page (best-effort)
                if ppdf is not None:
                    try:
                        page = ppdf.pages[page_index]
                        res = page.get("/Resources", None)
                        if res:
                            xobjs = res.get("/XObject", {})
                            # Try by name first
                            if xname and xname in xobjs:
                                xobj = xobjs[xname].get_object()
                                if isinstance(xobj, pikepdf.Stream):
                                    objnum, gennum = xobj.objgen
                                    colorspace = str(xobj.get("/ColorSpace", ""))
                                    f = xobj.get("/Filter", None)
                                    if isinstance(f, pikepdf.Array):
                                        filters = ",".join(str(i) for i in f)
                                    elif f is not None:
                                        filters = str(f)
                                    # SMask presence
                                    has_smask = xobj.get("/SMask", None) is not None
                            else:
                                # Brute: iterate and try to match by decoded bytes length and dimensions
                                for name, xref in xobjs.items():
                                    try:
                                        xobj = xref.get_object()
                                        if not isinstance(xobj, pikepdf.Stream):
                                            continue
                                        if xobj.get("/Subtype", None) != pikepdf.Name("/Image"):
                                            continue
                                        w = int(xobj.get("/Width", 0))
                                        h = int(xobj.get("/Height", 0))
                                        if width and height and (w != width or h != height):
                                            continue
                                        raw = xobj.read_bytes(decode_stream=True)
                                        if len(raw) == pixel_len and sha256_hex(raw) == pixel_sha:
                                            objnum, gennum = xobj.objgen
                                            xname = str(name)
                                            colorspace = str(xobj.get("/ColorSpace", ""))
                                            f = xobj.get("/Filter", None)
                                            if isinstance(f, pikepdf.Array):
                                                filters = ",".join(str(i) for i in f)
                                            elif f is not None:
                                                filters = str(f)
                                            has_smask = xobj.get("/SMask", None) is not None
                                            break
                                    except Exception:
                                        continue
                    except Exception:
                        pass

                out.append(ImgRec(
                    pdf_path=pdf_path,
                    page_index=page_index,
                    objnum=objnum,
                    gennum=gennum,
                    xobject_name=xname if isinstance(xname, str) else None,
                    width=width,
                    height=height,
                    bpc=bpc,
                    colorspace=colorspace,
                    filters=filters,
                    has_smask=has_smask,
                    pixel_len=pixel_len,
                    pixel_sha256=pixel_sha,
                    smask_len=smask_len,
                    smask_sha256=smask_sha,
                    source="inline" if (objnum is None) else "xobject",
                ))
    if ppdf is not None:
        ppdf.close()
    return out

def write_csv(rows: List[ImgRec], 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","source"
    ]
    with open(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]):
    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_page": lrec.page_index,
                    "left_obj": f"{lrec.objnum} {lrec.gennum}" if lrec.objnum is not None else "inline",
                    "left_dims": f"{lrec.width}x{lrec.height}",
                    "left_source": lrec.source,
                    "right_pdf": rrec.pdf_path,
                    "right_page": rrec.page_index,
                    "right_obj": f"{rrec.objnum} {rrec.gennum}" if rrec.objnum is not None else "inline",
                    "right_dims": f"{rrec.width}x{rrec.height}",
                    "right_source": rrec.source,
                    "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_page": lrec.page_index,
                    "left_obj": f"{lrec.objnum} {lrec.gennum}" if lrec.objnum is not None else "inline",
                    "left_dims": f"{lrec.width}x{lrec.height}",
                    "left_source": lrec.source,
                    "right_pdf": rrec.pdf_path,
                    "right_page": rrec.page_index,
                    "right_obj": f"{rrec.objnum} {rrec.gennum}" if rrec.objnum is not None else "inline",
                    "right_dims": f"{rrec.width}x{rrec.height}",
                    "right_source": rrec.source,
                    "sha256": h,
                })
    return pixel_matches, smask_matches

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

    write_csv(left, "left_images.csv")
    write_csv(right, "right_images.csv")
    print("[+] Wrote left_images.csv and right_images.csv")

    pixel_matches, smask_matches = compare_sets(left, right)

    with open("matches_pixels.csv", "w", newline="", encoding="utf-8") as f:
        w = csv.DictWriter(f, fieldnames=[
            "match_type","left_pdf","left_page","left_obj","left_dims","left_source",
            "right_pdf","right_page","right_obj","right_dims","right_source","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_page","left_obj","left_dims","left_source",
            "right_pdf","right_page","right_obj","right_dims","right_source","sha256"
        ])
        w.writeheader()
        for row in smask_matches:
            w.writerow(row)

    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:")
        print("   ", pixel_matches[0])

if __name__ == "__main__":
    main()