"""Compose reproducible inputs and measure real browser downloads; no removal code.

Requires Pillow and NumPy. Keep originals as supplied. Run:
  python restore_test.py compose --root . --alpha-map alpha-maps.js
  python restore_test.py measure --root .
manifest.json must list id, original, composited, cleaned and date for each case.
The cleaned files must be actual downloads from https://watermarkundo.com/.
"""
import argparse
import base64
import csv
import hashlib
import json
import re
from pathlib import Path
import numpy as np
from PIL import Image, __version__ as pillow_version


def sha(path):
    return hashlib.sha256(path.read_bytes()).hexdigest()


def pixels(path):
    with Image.open(path) as image:
        return np.array(image.convert('RGB'), dtype=np.uint8)


def geometry(width, height, manual_tier=None):
    tier, margin = (96, 64) if width > 1024 and height > 1024 else (48, 32)
    if manual_tier is not None:
        if manual_tier not in (48,96):
            raise ValueError('Unsupported manual tier')
        tier,margin=manual_tier,64 if manual_tier==96 else 32
    return tier, margin, width - margin - tier, height - margin - tier


def compare(original, other, x, y, tier):
    if original.shape != other.shape:
        raise ValueError('Dimensions changed')
    difference = np.abs(original.astype(np.int16) - other.astype(np.int16))
    changed = np.any(difference != 0, axis=2)
    inside = int(changed[y:y+tier, x:x+tier].sum())
    total = int(changed.sum())
    return {
        'max_channel_error': int(difference.max()),
        'absolute_error_sum_rgb': int(difference.sum()),
        'channel_count_rgb': int(difference.size),
        'mean_absolute_error_rgb': f'{difference.mean():.2f}',
        'mean_absolute_error_rgb_unrounded': float(difference.mean()),
        'mean_absolute_error_box_rgb': f'{difference[y:y+tier, x:x+tier].mean():.2f}',
        'changed_pixels_inside_box': inside,
        'changed_pixels_outside_box': total - inside,
        'changed_pixels_total': total,
    }


def compose(root, manifest, map_path):
    text = map_path.read_text(encoding='utf-8')
    maps = dict(re.findall(r"'(48|96)'\s*:\s*'([A-Za-z0-9+/=]+)'", text))
    for case in manifest['cases']:
        source_path = root / case['original']
        original = pixels(source_path)
        height, width, _ = original.shape
        tier, margin, x, y = geometry(width, height,case.get('manual_tier'))
        raw = np.frombuffer(base64.b64decode(maps[str(tier)]), dtype='<f4').reshape(tier,tier).astype(np.float64)
        magnitude = np.abs(raw)
        # Same support/noise gate and signed logo convention as current app.js.
        active = np.maximum(0.0, magnitude - 3/255) >= 0.002
        alpha = np.where(active, np.minimum(magnitude, 0.99), 0.0)
        logo = np.where(raw < 0, 0.0, 255.0)
        patch = original[y:y+tier,x:x+tier].astype(np.float64)
        blend = (1-alpha[...,None])*patch + alpha[...,None]*logo[...,None]
        composited = original.copy()
        composited[y:y+tier,x:x+tier] = np.clip(np.floor(blend+0.5),0,255).astype(np.uint8)
        output_path = root / case['composited']
        Image.fromarray(composited).save(output_path, format='PNG', compress_level=6)
        with Image.open(source_path) as image:
            case.update(width=width,height=height,format=image.format)
        case.update(tier=tier,auto_tier=geometry(width,height)[0],selection_mode='Manual' if case.get('manual_tier') else 'Auto',margin=margin,box_x=x,box_y=y,box_width=tier,box_height=tier,
                    strength='1.00',original_bytes=source_path.stat().st_size,
                    input_bytes=output_path.stat().st_size,original_sha256=sha(source_path),
                    composited_sha256=sha(output_path),
                    input_vs_ground_truth=compare(original,composited,x,y,tier))
    manifest['method'] = {
        'baseline': 'RGB pixels decoded by Pillow from owner-supplied JPEG originals, preserved byte for byte',
        'composition': 'quantize to uint8 by floor(value+0.5); value=(1-alpha)*ground_truth+alpha*logo',
        'alpha': 'Exact signed Float32 map from the unchanged site alpha-maps.js; magnitude clipped at 0.99; logo=0 if signed map <0 else 255; ignore support where max(0,abs(raw)-3/255)<0.002; gain=1.00',
        'alpha_map_sha256': sha(map_path),
        'geometry': 'Auto: width>1024 AND height>1024: size96 margin64; otherwise size48 margin32. Owner-approved fourth case manually uses size96 margin64 on1200x896; it does not validate Auto96. x=width-margin-size; y=height-margin-size; left/top inclusive, right/bottom exclusive',
        'measurement': 'RGB, scale0..255; max over all channels; MAE=sum absolute RGB differences/(width*height*3); count pixel once if any RGB channel differs; compare cleaned with ground truth, not input',
        'versions': {'Pillow': pillow_version,'NumPy':np.__version__},
        'cleanup': 'Not implemented in this script. Cleaned files must be downloaded from live tool: first three Auto48, fourth manually selected96; strength1.00, offsets0, overlay off.',
        'limitations': 'Known-map composited inputs, not native Gemini watermarked downloads; not success rate or evidence for all images; SynthID, metadata and mobile processing not tested.',
    }
    return manifest


def measure(root, manifest):
    complete=[]
    for case in manifest['cases']:
        path=root/case['cleaned']
        if not path.exists():
            continue
        if case.get('live_tool_readout') is None:
            raise ValueError('Record live-tool readout before measuring '+case['id'])
        original=pixels(root/case['original'])
        result=pixels(path)
        case['cleaned_vs_ground_truth']=compare(original,result,case['box_x'],case['box_y'],case['tier'])
        case['cleaned_bytes']=path.stat().st_size
        case['cleaned_sha256']=sha(path)
        row={key:case[key] for key in ['id','width','height','selection_mode','auto_tier','tier','margin','box_x','box_y','original_bytes','input_bytes','cleaned_bytes','date']}
        row.update(case['cleaned_vs_ground_truth'])
        row['input_changed_pixels_vs_ground_truth']=case['input_vs_ground_truth']['changed_pixels_total']
        complete.append(row)
    if complete:
        with (root/'measurements.csv').open('w',encoding='utf-8',newline='') as file:
            writer=csv.DictWriter(file,fieldnames=list(complete[0]))
            writer.writeheader();writer.writerows(complete)
    return manifest


if __name__=='__main__':
    parser=argparse.ArgumentParser()
    parser.add_argument('action',choices=['compose','measure'])
    parser.add_argument('--root',type=Path,required=True)
    parser.add_argument('--alpha-map',type=Path)
    args=parser.parse_args()
    root=args.root.resolve()
    path=root/'manifest.json'
    manifest=json.loads(path.read_text(encoding='utf-8-sig'))
    if args.action=='compose':
        manifest=compose(root,manifest,args.alpha_map or root/'alpha-maps.js')
    else:
        manifest=measure(root,manifest)
    path.write_text(json.dumps(manifest,ensure_ascii=False,indent=2),encoding='utf-8')
    print(json.dumps(manifest,ensure_ascii=False,indent=2))