#!/usr/bin/env python3
"""Compare a real ICCBased PDF through the sandboxed worker at three scales."""
from pathlib import Path
from io import BytesIO
import argparse
import hashlib
import json
import math
import os
import subprocess
from PIL import Image, ImageChops, ImageDraw, ImageStat, ImageCms

ROOT=Path(__file__).resolve().parent.parent
parser=argparse.ArgumentParser(description=__doc__)
parser.add_argument("--build-dir",type=Path,default=ROOT/os.environ.get("DOCVIEW_BUILD_DIR","build"))
parser.add_argument("--output",type=Path,default=ROOT/"tests/results/pdf-icc")
args=parser.parse_args()
from pypdf import PdfReader
BUILD=args.build_dir.resolve()
CORPUS=ROOT/"tests/fixtures/pdf/extended"
OUT=args.output.resolve(); OUT.mkdir(parents=True,exist_ok=True)
manifest=json.loads((CORPUS/"manifest.json").read_text())
spec=next(item for item in manifest["documents"] if item["file"]=="icc-linear-rgb.pdf")
source=CORPUS/spec["file"]; before=hashlib.sha256(source.read_bytes()).hexdigest()
assert before==spec["sha256"]
profile=(CORPUS/"linear-srgb.icc").read_bytes()
assert hashlib.sha256(profile).hexdigest()==manifest["color_profile"]["sha256"]
reader=PdfReader(source)
color_space=reader.pages[0]["/Resources"]["/ColorSpace"]["/Linear"]
assert str(color_space[0])=="/ICCBased"
assert color_space[1].get_object().get_data()==profile
image=reader.pages[0]["/Resources"]["/XObject"]["/Sample"].get_object()
assert str(image["/ColorSpace"][0])=="/ICCBased"
assert image["/ColorSpace"][1].get_object().get_data()==profile
cms_profile=ImageCms.ImageCmsProfile(BytesIO(profile))
report={"scope":"Self-authored ICCBased RGB matrix/TRC vector and raster at actual worker render scales; not full ICC/CMYK/printing acceptance",
    "fixture_sha256_before":before,"profile":manifest["color_profile"],
    "worker_sha256":hashlib.sha256((BUILD/"docview-pdf-worker").read_bytes()).hexdigest(),
    "profile_copyright_readback":ImageCms.getProfileCopyright(cms_profile).strip(),
    "profile_description_readback":ImageCms.getProfileDescription(cms_profile).strip(),
    "profile_embedded_in_vector_and_image":True,"transport":"production sandboxed PDFWorker, WorkerProcess, FontBroker, tile IPC",
    "poppler_version":subprocess.run(["/usr/bin/pdftoppm","-v"],capture_output=True,text=True).stderr.splitlines()[0],
    "sha256":{str(path.relative_to(ROOT)):hashlib.sha256(path.read_bytes()).hexdigest() for path in
        [BUILD/"pdf-worker-evidence",BUILD/"docview-pdf-worker",BUILD/"libdocview_pdf.a",BUILD/"libdocview_font_broker.a",ROOT/"tests/render_pdf_worker.cpp",ROOT/"tests/compare_pdf_icc.py",ROOT/"tests/fixtures/pdf/icc_generator.py"]},
    "reference":"Analytic sRGB transfer function of stated linear RGB values; independent Poppler raster; no human-approved golden",
    "scales":[]}
overview=[]; success=True
for scale in [.5,1.,4.]:
    folder=OUT/("scale-"+str(scale).replace(".","_")); (folder/"pdfium").mkdir(parents=True,exist_ok=True)
    subprocess.run([str(BUILD/"pdf-worker-evidence"),str(source),str(folder/"pdfium"),str(scale)],check=True)
    subprocess.run(["/usr/bin/pdftoppm","-singlefile","-cropbox","-r",str(72*scale),"-png",str(source),str(folder/"poppler")],check=True)
    a=Image.open(folder/"pdfium/page-1.png").convert("RGB"); b=Image.open(folder/"poppler.png").convert("RGB")
    meta=json.loads((folder/"pdfium/metadata.json").read_text()); assert meta["renderScale"]==scale
    size=[math.ceil(v*scale) for v in spec["expected"]["page_sizes_pt"][0]]
    entry={"scale":scale,"dpi":72*scale,"expected_size":size,"pdfium_size":list(a.size),"poppler_size":list(b.size),
        "dimensions_match_source":list(a.size)==size==list(b.size),"probes":[]}
    success=success and entry["dimensions_match_source"]
    for probe in spec["expected"]["icc_probes"]:
        x,y=probe["point_pt"]; point=(round(x*scale),round((350-y)*scale))
        actual_a,actual_b=a.getpixel(point),b.getpixel(point)
        expected=probe["expected_rgb"]; tolerance=spec["expected"]["probe_tolerance_rgb"]
        check=dict(probe,point_px=point,pdfium_rgb=actual_a,poppler_rgb=actual_b,
            both_match_analytic_color=all(abs(actual[k]-expected[k])<=tolerance for actual in (actual_a,actual_b) for k in range(3)))
        entry["probes"].append(check);success=success and check["both_match_analytic_color"]
    diff=ImageChops.difference(a,b);diff.save(folder/"difference.png")
    entry["mean_absolute_channel_error"]=sum(ImageStat.Stat(diff).mean)/3
    entry["icc_midgray_differs_from_device_control"]=all(abs(entry["probes"][6][which][0]-entry["probes"][7][which][0])>40 for which in ("pdfium_rgb","poppler_rgb"))
    success=success and entry["icc_midgray_differs_from_device_control"]
    sheet=Image.new("RGB",(1500,350),"#ededed");draw=ImageDraw.Draw(sheet)
    for col,(name,img) in enumerate([("PDFium / production worker",a),("Poppler / independent",b),("Absolute RGB difference",diff)]):
        draw.text((col*500+10,12),f"{scale:g}x {name}",fill="black");thumb=img.copy();thumb.thumbnail((488,300))
        sheet.paste(thumb,(col*500+(500-thumb.width)//2,42))
    sheet.save(folder/"comparison.png");overview.append(sheet);report["scales"].append(entry)
report["fixture_sha256_after"]=hashlib.sha256(source.read_bytes()).hexdigest()
report["original_unchanged"]=report["fixture_sha256_after"]==before;success=success and report["original_unchanged"]
report["all_source_assertions_pass"]=success
canvas=Image.new("RGB",(1500,len(overview)*350),"white")
for i,sheet in enumerate(overview):canvas.paste(sheet,(0,i*350))
canvas.save(OUT/"overview.png")
(OUT/"comparison.json").write_text(json.dumps(report,ensure_ascii=False,indent=2)+"\n")
print(json.dumps({"all_source_assertions_pass":success,"scales":len(report["scales"]),"probe_count":sum(len(x["probes"]) for x in report["scales"]),"original_unchanged":report["original_unchanged"],"profile_copyright":report["profile_copyright_readback"]},indent=2))
if not success: raise SystemExit(f"ICC source assertion failed; see {OUT / 'comparison.json'}")