#!/usr/bin/env python3
"""Validate a known CMYK CLUT through the production isolated worker and Poppler."""
import argparse
import hashlib
from io import BytesIO
import json
import math
import os
from pathlib import Path
import platform
import struct
import subprocess

from PIL import Image, ImageChops, ImageCms, ImageDraw, ImageStat
from pypdf import PdfReader

ROOT = Path(__file__).resolve().parents[2]


def sha(path):
    with path.open('rb') as stream:
        return hashlib.file_digest(stream, 'sha256').hexdigest()


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--build-dir', required=True, type=Path)
    parser.add_argument('--output', required=True, type=Path)
    args = parser.parse_args()
    build = args.build_dir.resolve(); output = args.output.resolve()
    output.mkdir(parents=True, exist_ok=False)
    corpus = ROOT / 'tests/fixtures/pdf/cmyk-lut'
    spec = json.loads((corpus / 'manifest.json').read_text())
    source = corpus / spec['file']; icc_file = corpus / spec['profileFile']
    assert sha(source) == spec['sha256'] and sha(icc_file) == spec['profileSha256']
    assert sha(ROOT / 'tests/fixtures/pdf/generate_cmyk_lut.py') == spec['generatorSha256']
    binaries = {name: sha(build / name) for name in ('pdf-worker-evidence', 'docview-pdf-worker')}
    profile = icc_file.read_bytes()
    assert len(profile) == spec['profileBytes'] == struct.unpack_from('>I', profile)[0]
    assert profile[8:24] == bytes.fromhex('02100000') + b'scnrCMYKLab ' and profile[36:40] == b'acsp'
    tags = {}
    for i in range(struct.unpack_from('>I', profile, 128)[0]):
        sig, offset, size = struct.unpack_from('>4sII', profile, 132+12*i)
        assert offset % 4 == 0 and offset+size <= len(profile) and sig not in tags
        tags[sig] = profile[offset:offset+size]
    assert set(tags) == {b'desc', b'cprt', b'wtpt', b'A2B0'}
    lut = tags[b'A2B0']
    assert len(lut) == 176 and lut[:12] == b'mft2' + bytes(4) + bytes((4,3,2,0))
    assert struct.unpack_from('>HH', lut, 48) == (2,2)
    # Decode and check all 16 four-dimensional CLUT vertices independently.
    for index in range(16):
        c, m, y, k = [(index >> shift) & 1 for shift in (3,2,1,0)]
        encoded = struct.unpack_from('>3H', lut, 68+index*6)
        lab = [encoded[0]*100/65280, encoded[1]/256-128, encoded[2]/256-128]
        expected = [100-12*c-15*m-10*y-35*k, 18*m-15*c, 20*y-12*c]
        assert max(abs(a-b) for a,b in zip(lab,expected)) < .002
    pdf = PdfReader(source)
    assert len(pdf.pages) == 1 and list(pdf.pages[0].mediabox)[2:] == spec['pageSizePt']
    resources = pdf.pages[0]['/Resources']
    vector_space = resources['/ColorSpace']['/TestCMYK']
    image_object = resources['/XObject']['/Sample'].get_object()
    for space in (vector_space, image_object['/ColorSpace']):
        assert space[0] == '/ICCBased' and space[1].get_object()['/N'] == 4
        assert space[1].get_object()['/Alternate'] == '/DeviceCMYK'
        assert space[1].get_object().get_data() == profile
    assert image_object['/Width'] == 8 and image_object['/Height'] == 1 and image_object['/BitsPerComponent'] == 8
    pixel_values = bytes(round(c*255) for p in spec['probes'] if p['kind'] == 'icc-image' for c in p['components'])
    assert image_object.get_data() == pixel_values
    icc = ImageCms.ImageCmsProfile(BytesIO(profile))
    transform = ImageCms.buildTransformFromOpenProfiles(icc, ImageCms.createProfile('sRGB'), 'CMYK', 'RGB', renderingIntent=1)
    cmm_probes = []
    for probe in spec['probes']:
        if probe['kind'] != 'icc-image': continue
        pixel = Image.new('CMYK',(1,1),tuple(round(c*255) for c in probe['components']))
        actual = ImageCms.applyTransform(pixel, transform).getpixel((0,0))
        error = max(abs(a-b) for a,b in zip(actual,probe['expectedRgb']))
        cmm_probes.append({'sample':probe['sample'], 'expectedRgb':probe['expectedRgb'], 'actualRgb':actual,
                           'maxChannelError':error, 'success':error <= spec['tolerance8Bit']})
    report = {'success':False, 'osRelease':platform.freedesktop_os_release(), 'fixture':spec,
              'python':platform.python_version(), 'pillowLcmsVersion':ImageCms.core.littlecms_version,
              'runnerSha256':sha(Path(__file__)), 'binaries':binaries,
              'rendererVersion':subprocess.run(['/usr/bin/pdftoppm','-v'],capture_output=True,text=True,check=True).stderr.splitlines()[0],
              'embeddedProfileAndImageBytesVerified':True, 'clutVerticesVerified':16, 'cmmProbes':cmm_probes,
              'scope':'Synthetic ICC v2 CMYK input -> Lab four-dimensional CLUT, vectors and images. '
                      'Production isolated worker at three scales and independent Poppler. '
                      'The CMM probe uses Pillow/LittleCMS and is not a third independent PDF renderer. '
                      'No calibrated press/display, arbitrary ICC conformance or human-approved golden.',
              'commands':[], 'scales':[]}

    def run(name, command):
        with (output / (name+'.log')).open('w') as log:
            result = subprocess.run(command, stdout=log, stderr=subprocess.STDOUT, timeout=180)
        report['commands'].append({'name':name, 'command':command, 'exitCode':result.returncode,
                                   'logSha256':sha(output / (name+'.log'))})
        assert result.returncode == 0, name

    try:
        assert all(p['success'] for p in cmm_probes), 'Profile differs from the analytic model'
        run('qpdf-check',['/usr/bin/qpdf','--check',str(source)])
        sheets=[]
        for scale in (.5, 1., 4.):
            name = 'scale-'+str(scale).replace('.','_'); folder=output/name
            (folder/'pdfium').mkdir(parents=True)
            run(name+'-pdfium',[str(build/'pdf-worker-evidence'),str(source),str(folder/'pdfium'),str(scale)])
            run(name+'-poppler',['/usr/bin/pdftoppm','-singlefile','-cropbox','-r',str(72*scale),'-png',str(source),str(folder/'poppler')])
            meta=json.loads((folder/'pdfium/metadata.json').read_text())
            assert meta['renderScale']==scale and 'sandbox enabled' in meta['transport']
            images={'pdfium':Image.open(folder/'pdfium/page-1.png').convert('RGB'),
                    'poppler':Image.open(folder/'poppler.png').convert('RGB')}
            dimensions=[math.ceil(v*scale) for v in spec['pageSizePt']]
            assert all(list(im.size)==dimensions for im in images.values())
            row={'scale':scale, 'dpi':72*scale, 'dimensions':dimensions, 'probes':[], 'success':True}
            for probe in spec['probes']:
                x,y=probe['pointPt']; point=(round(x*scale),round((spec['pageSizePt'][1]-y)*scale))
                colors={key:list(im.getpixel(point)) for key,im in images.items()}
                errors={key:max(abs(a-b) for a,b in zip(color,probe['expectedRgb'])) for key,color in colors.items()} if 'expectedRgb' in probe else {}
                matched=all(error<=spec['tolerance8Bit'] for error in errors.values())
                row['probes'].append({**probe,'pointPx':point,'actualRgb':colors,'maxChannelError':errors,
                                       'checkedAgainstAnalyticColor':bool(errors),'success':matched})
                row['success'] &= matched
            # A profile ignored in favor of its DeviceCMYK Alternate cannot pass.
            black={p['kind']:p for p in row['probes'] if p['sample']==4}
            differences={key:min(black['icc-vector']['actualRgb'][key][i]-black['device-control']['actualRgb'][key][i]
                                 for i in range(3)) for key in images}
            row['blackProfileMinusAlternateMinChannel']=differences
            row['profileEffectDetected']=all(value>=80 for value in differences.values())
            row['success'] &= row['profileEffectDetected']
            diff=ImageChops.difference(images['pdfium'],images['poppler']); diff.save(folder/'difference.png')
            row['meanAbsoluteChannelError']=sum(ImageStat.Stat(diff).mean)/3
            sheet=Image.new('RGB',(1500,340),'#eeeeee'); draw=ImageDraw.Draw(sheet)
            for col,(label,im) in enumerate([('PDFium - isolated worker',images['pdfium']),('Poppler - independent',images['poppler']),('Absolute RGB difference',diff)]):
                draw.text((col*500+10,10),f'{scale:g}x / {label}',fill='black')
                thumb=im.copy(); thumb.thumbnail((488,295)); sheet.paste(thumb,(col*500+(500-thumb.width)//2,35))
            sheet.save(folder/'comparison.png'); sheets.append(sheet); report['scales'].append(row)
        overview=Image.new('RGB',(1500,340*len(sheets)),'white')
        for i,sheet in enumerate(sheets): overview.paste(sheet,(0,340*i))
        overview.save(output/'overview.png')
        report['originalUnchanged']=sha(source)==spec['sha256'] and sha(icc_file)==spec['profileSha256']
        report['binariesUnchanged']=all(sha(build/name)==digest for name,digest in binaries.items())
        report['success']=all(row['success'] for row in report['scales']) and report['originalUnchanged'] and report['binariesUnchanged']
    finally:
        report['evidenceSha256']={str(path.relative_to(output)):sha(path) for path in sorted(output.rglob('*')) if path.is_file()}
        (output/'report.json').write_text(json.dumps(report,indent=2)+'\n')
    print(json.dumps({'success':report['success'],'scales':len(report['scales']),
                      'analyticImageAndVectorProbes':sum(p['checkedAgainstAnalyticColor'] for row in report['scales'] for p in row['probes'])}))
    if not report['success']: raise SystemExit('Color validation failed; inspect the retained report')


if __name__=='__main__':
    main()