"""Validate supplementary editing pairs against the release's model identities."""
import json
import hashlib
import zipfile
from pathlib import Path
from PIL import Image

from scripts.integrity import sha256, validate_pair
from scripts.report import load_records


def validate_editing(evaluation):
    evaluation = Path(evaluation)
    root = evaluation/'editing'
    summary = json.loads((root/'summary.json').read_text(encoding='utf-8'))
    cases = json.loads((root/'cases.json').read_text(encoding='utf-8'))
    # Hub text files use LF; preserve the exact measured Windows case-file bytes separately.
    with zipfile.ZipFile(root/'measured-cases.zip') as archive:
        measured_cases = archive.read('editing.json')
    if json.loads(measured_cases) != cases:
        raise ValueError('Published case definitions differ from the measured case file')
    expected = {(case['id'], 42) for case in cases}
    if len(expected) != 3:
        raise ValueError('Expected three declared editing cases')
    records = {p: load_records(root/p) for p in ('bf16', 'int8')}
    env = {p: json.loads((root/p/'environment.json').read_text()) for p in records}
    for p in records:
        core = json.loads((evaluation/p/'environment.json').read_text())
        if env[p]['model_identity']['fingerprint'] != core['model_identity']['fingerprint']:
            raise ValueError('Editing model differs from main suite')
        if records[p].keys() != expected or env[p]['before_load_memory']['allocated_bytes'] != 0:
            raise ValueError('Incomplete editing set or nonzero startup allocation')
        if not env[p]['warmup'] or not env[p]['vae_tiling'] or not env[p]['offload_aux_fix']:
            raise ValueError('Missing warmup, VAE tiling or auxiliary offload')
        if env[p]['cases_sha256'] != hashlib.sha256(measured_cases).hexdigest():
            raise ValueError('Declared editing cases changed')
    for key in ('gpu', 'cuda', 'packages', 'offload', 'generator_device', 'source_sha256', 'vae_tiles'):
        if env['bf16'][key] != env['int8'][key]:
            raise ValueError('Editing runtime mismatch: '+key)
    source_root = evaluation.parent/'scripts'
    if source_root.exists():
        for name, digest in env['bf16']['source_sha256'].items():
            if sha256(source_root/name) != digest:
                raise ValueError('Published editing implementation differs from measured code')
    measures = []
    for case in cases:
        a, b = (records[p][(case['id'], 42)] for p in ('bf16', 'int8'))
        validate_pair(a, b)
        input_hash = sha256(root/'inputs'/case['input_file'])
        source = Path(case['source'])
        if source.parts[0] != 'evaluation' or '..' in source.parts:
            raise ValueError('Invalid source-image provenance path')
        if sha256(evaluation/Path(*source.parts[1:])) != input_hash:
            raise ValueError('Input differs from its documented BF16 source image')
        for row in (a, b):
            if (row['prompt'] != case['prompt'] or row['input_sha256'] != input_hash
                    or row['input_sha256'] != case['input_sha256']):
                raise ValueError('Editing prompt/input mismatch')
            if (row['width'], row['height'], row['output_resolution'], row['steps'],
                row['cfg'], row['kv_cache'], row['vae_tiling'], row['offload']) != (
                    2048, 2048, 2048, 40, 1.0, True, True, 'model'):
                raise ValueError('Unexpected editing settings')
        for p in records:
            row = records[p][(case['id'], 42)]
            with Image.open(root/p/row['image']) as image:
                if list(image.size) != row['actual_size'] or image.size != (2048, 2048):
                    raise ValueError('Editing PNG dimensions differ from recorded dimensions')
                if image.mode != row['image_mode']:
                    raise ValueError('Editing PNG mode differs from recorded mode')
        measures.append(dict(case_id=case['id'], bf16_seconds=a['seconds'], int8_seconds=b['seconds'],
                             bf16_peak_gib=a['peak_allocated_bytes']/2**30,
                             int8_peak_gib=b['peak_allocated_bytes']/2**30))
    if measures != summary['measurements']:
        raise ValueError('Editing summary differs from measured records')
    if [c['id'] for c in summary['cases']] != ['sweater_red'] + [c['id'] for c in cases]:
        raise ValueError('Editing gallery does not include every declared case')
    for item, case in zip(summary['cases'][1:], cases):
        if item['prompt'] != case['prompt'] or item['input_sha256'] != case['input_sha256']:
            raise ValueError('Gallery prompt/input differs from declared case')
        for p in records:
            record = records[p][(case['id'], 42)]
            if item[p+'_sha256'] != record['image_sha256']:
                raise ValueError('Gallery output differs from measured record')
    native = json.loads((evaluation/'audit/summary.json').read_text())['native']
    if summary['cases'][0]['input_sha256'] != sha256(evaluation/'bf16/portrait-s42.png'):
        raise ValueError('Existing sweater input differs from the BF16 portrait')
    for p in records:
        if summary['cases'][0][p+'_sha256'] != native[p]['image_sha256']:
            raise ValueError('Existing sweater image differs from retained evidence')
    for item in summary['cases']:
        for key in ('input', 'bf16', 'int8'):
            path = Path(item[key])
            if path.parts[0] != 'evaluation' or '..' in path.parts:
                raise ValueError('Invalid editing gallery path')
            if sha256(evaluation/Path(*path.parts[1:])) != item[key+'_sha256']:
                raise ValueError('Editing gallery hash mismatch')
    isolation = [json.loads(line) for line in (root/'process-isolation.jsonl').read_text().splitlines()]
    if [row['precision'] for row in isolation] != ['bf16', 'int8']:
        raise ValueError('Unexpected editing process order')
    for row in isolation:
        if row['exit_code'] != 0 or row['inference_pid'] != env[row['precision']]['pid']:
            raise ValueError('Editing process did not finish successfully')
        after_pids = {line.split(',')[0].strip() for line in row['gpu_processes_after'].splitlines()[1:]}
        if str(row['inference_pid']) in after_pids:
            raise ValueError('Editing process remains resident after exit')
    return summary