#!/usr/bin/env python3
"""Standalone Re:InterHand Test40 template / verify / pack. Python 3.10+, standard library only.

Uses the ORIGINAL inputs.json from the selected frozen dataset package.
Never modifies that package. Templates are marked as fixtures by default.
"""
import argparse
import hashlib
import json
import math
from pathlib import Path
import zipfile

SCHEMA = 'reinterhand-pose-prediction-v1'
CAMERAS = ('ego',)
SIDES = ('left', 'right')
MAX_CASE_BYTES = 32 * 1024**2
MAX_TOTAL_BYTES = 512 * 1024**2
MAX_METADATA_BYTES = 128 * 1024


def read(path):
    return json.loads(path.read_text())


def write(path, value):
    path.write_text(json.dumps(value, ensure_ascii=False, indent=2, allow_nan=False) + '\n')


def load_inputs(inputs):
    source = read(inputs)
    expected = {'schema_version': 'reinterhand-inputs-v1',
        'protocol_id': 'reinterhand-ego-mano21-pose-v1',
        'dataset_id': 'reinterhand-ego-envmap-per-segment-mano21-test40-v1',
        'split_id': 'test40-v1', 'preprocessing_id': 'reinterhand-ego-native-rgb-no-transform-v1',
        'coordinate_frame': 'camera:ego:native', 'unit': 'm', 'camera_ids': ['ego']}
    if any(source.get(k) != v for k, v in expected.items()):
        raise ValueError('expected original Test40 inputs.json')
    clips = source['clips']
    if len(clips) != 40 or len({c['clip_id'] for c in clips}) != 40 or sum(c['frame_count'] for c in clips) != 20664:
        raise ValueError('expected 40 full-length cases / 20664 frames')
    import re
    for clip in clips:
        if not re.fullmatch(r'reih-(?:rns217|nkc880|qzx685)-[a-z0-9_]+-\d{2}', clip['clip_id']):
            raise ValueError('invalid case ID')
        if type(clip['frame_count']) is not int or not 99 <= clip['frame_count'] <= 1761:
            raise ValueError('invalid full-length frame count')
        if clip['cameras'] != {'ego': {'width': 1024, 'height': 1024}}:
            raise ValueError('expected native Ego image dimensions')
    if len(source['joint_names']) != 21 or len(set(source['joint_names'])) != 21:
        raise ValueError('expected 21 unique joint names')
    return source


def verify(inputs, submission):
    if submission.is_symlink() or (submission / 'submission.json').is_symlink():
        raise ValueError('symlinks are not supported')
    if (submission / 'submission.json').stat().st_size > MAX_METADATA_BYTES:
        raise ValueError('submission.json exceeds 128 KiB')
    source, meta = load_inputs(inputs), read(submission / 'submission.json')
    if not isinstance(meta, dict):
        raise ValueError('submission.json must contain an object')
    if meta.get('submission_kind', 'model_output') not in ('model_output', 'gt_derived_fixture'):
        raise ValueError('submission_kind must be model_output or gt_derived_fixture')
    if meta['schema_version'] != SCHEMA:
        raise ValueError('expected Re:InterHand prediction metadata')
    kind, cameras = meta['prediction_type'], meta['inference_camera_ids']
    if kind not in ('3d', '2d') or not cameras or len(set(cameras)) != len(cameras) or set(cameras) - set(CAMERAS):
        raise ValueError('invalid prediction_type/inference_camera_ids')
    expected = {k: source[k] for k in ('protocol_id', 'dataset_id', 'preprocessing_id', 'split_id')}
    expected.update(input_manifest_sha256=hashlib.sha256(inputs.read_bytes()).hexdigest(),
                    coordinate_frame='image:ego:native' if kind == '2d' else source['coordinate_frame'],
                    unit='px' if kind == '2d' else 'm')
    if any(meta.get(k) != v for k, v in expected.items()):
        raise ValueError('input identity or coordinates differ; use original inputs.json')
    if len(meta['joint_names']) != 21 or set(meta['joint_names']) != set(source['joint_names']):
        raise ValueError('declare 21 unique semantic joints')
    for key in ('pipeline_version', 'prediction_adapter_id'):
        if not isinstance(meta.get(key), str) or not meta[key].strip():
            raise ValueError(key + ' is required')
    clips = {c['clip_id']: c for c in source['clips']}
    ids = meta['case_ids']
    if not 1 <= len(ids) <= 40 or len(set(ids)) != len(ids) or set(ids) - clips.keys():
        raise ValueError('invalid case_ids')
    expected_files = {'submission.json'} | {'predictions/' + cid + '.jsonl' for cid in ids}
    if submission.is_symlink() or any(p.is_symlink() for p in submission.rglob('*')):
        raise ValueError('symlinks are not supported')
    if {p.relative_to(submission).as_posix() for p in submission.rglob('*') if p.is_file()} != expected_files:
        raise ValueError('unexpected or missing submission files')
    if sum((submission / name).stat().st_size for name in expected_files) > MAX_TOTAL_BYTES:
        raise ValueError('submission exceeds 512 MiB uncompressed limit')
    for cid in ids:
        if (submission / 'predictions' / (cid + '.jsonl')).stat().st_size > MAX_CASE_BYTES:
            raise ValueError('prediction case exceeds 32 MiB')
        rows = [json.loads(line) for line in (submission / 'predictions' / (cid + '.jsonl')).read_text().splitlines()]
        if len(rows) != clips[cid]['frame_count'] or any(type(r.get('frame_id')) is not int for r in rows) or sorted(r['frame_id'] for r in rows) != list(range(clips[cid]['frame_count'])):
            raise ValueError(cid + ': require all N unique frames from inputs.json')
        for row in rows:
            if row.get('clip_id') != cid:
                raise ValueError('clip identity differs')
            if kind == '2d':
                if set(row.get('views', {})) != set(cameras) or any(s in row for s in SIDES):
                    raise ValueError('2D must have exactly the declared views, no 3D slots')
                views = [row['views'][c] for c in cameras]
                for c, view in zip(cameras, views):
                    if any(type(view.get(k)) is not int or view[k] != clips[cid]['cameras'][c][k] for k in ('width', 'height')):
                        raise ValueError('2D dimensions must match native Ego image')
            else:
                if 'views' in row:
                    raise ValueError('3D must not contain 2D views')
                views = [row]
            field, size = ('points_2d', 2) if kind == '2d' else ('points_3d', 3)
            for view in views:
                for side in SIDES:
                    hand = view[side]
                    if any(k in hand for k in (('points_3d', 'uv_upright') if kind == '2d' else ('points_2d', 'uv_upright'))):
                        raise ValueError('mixed modality is not supported')
                    if hand.get('status') == 'missing' and field in hand and hand[field] is None:
                        continue
                    points = hand.get(field)
                    if hand.get('status') != 'ok' or not isinstance(points, list) or len(points) != 21 or any(
                        not isinstance(point, list) or len(point) != size or any(type(x) not in (float, int) or not math.isfinite(x) or (kind == '3d' and not -1e6 <= x <= 1e6) for x in point) for point in points):
                        raise ValueError('ok requires 21 finite points; otherwise explicitly mark missing')
    return meta


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('command', choices=('template', 'verify', 'pack'))
    parser.add_argument('--inputs', type=Path, required=True)
    parser.add_argument('--submission', type=Path, required=True)
    parser.add_argument('--type', choices=('3d', '2d'), default='3d')
    parser.add_argument('--cameras', nargs='+', choices=CAMERAS, default=list(CAMERAS))
    parser.add_argument('--case-id', action='append')
    parser.add_argument('--output', type=Path)
    parser.add_argument('--model-output', action='store_true', help='Only for real model output; never use for GT-derived fixtures')
    args = parser.parse_args()
    if args.command == 'template':
        source = load_inputs(args.inputs)
        ids = sorted(set(args.case_id or [c['clip_id'] for c in source['clips']]))
        if args.submission.exists() or set(ids) - {c['clip_id'] for c in source['clips']} or len(set(args.cameras)) != len(args.cameras):
            parser.error('new submission directory, valid cases and unique cameras required')
        (args.submission / 'predictions').mkdir(parents=True)
        meta = {k: source[k] for k in ('protocol_id', 'dataset_id', 'preprocessing_id', 'split_id', 'joint_names')}
        meta.update(schema_version=SCHEMA, prediction_type=args.type, inference_camera_ids=[c for c in CAMERAS if c in args.cameras],
                    case_ids=ids, input_manifest_sha256=hashlib.sha256(args.inputs.read_bytes()).hexdigest(),
                    coordinate_frame='image:ego:native' if args.type == '2d' else source['coordinate_frame'], unit='px' if args.type == '2d' else 'm',
                    prediction_adapter_id='replace-with-your-adapter-version', pipeline_version='replace-with-your-pipeline-version', submission_kind='gt_derived_fixture')
        write(args.submission / 'submission.json', meta)
        for clip in source['clips']:
            if clip['clip_id'] not in ids:
                continue
            rows = []
            for fid in range(clip['frame_count']):
                row = {'clip_id': clip['clip_id'], 'frame_id': fid}
                hands = {s: {'status': 'missing', 'points_2d' if args.type == '2d' else 'points_3d': None} for s in SIDES}
                if args.type == '2d':
                    row['views'] = {c: {**{k: clip['cameras'][c][k] for k in ('width', 'height')}, **hands} for c in meta['inference_camera_ids']}
                else:
                    row.update(hands)
                rows.append(row)
            (args.submission / 'predictions' / (clip['clip_id'] + '.jsonl')).write_text(''.join(json.dumps(r) + '\n' for r in rows))
    meta = verify(args.inputs, args.submission)
    if args.command == 'pack':
        if not args.output or args.output.exists():
            parser.error('provide a new --output ZIP')
        if args.model_output:
            meta['submission_kind'] = 'model_output'
        with zipfile.ZipFile(args.output, 'x', zipfile.ZIP_DEFLATED) as archive:
            archive.writestr('submission.json', json.dumps(meta, ensure_ascii=False, allow_nan=False))
            for cid in meta['case_ids']:
                path = args.submission / 'predictions' / (cid + '.jsonl')
                archive.write(path, 'predictions/' + path.name)
    print(json.dumps({'status': 'ok', 'case_count': len(meta['case_ids']), 'prediction_type': meta['prediction_type'], 'inference_camera_ids': meta['inference_camera_ids']}))


if __name__ == '__main__':
    main()