#!/usr/bin/env python3
"""Standalone HOT3D v4 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 = 'hot3d-pose-prediction-v4'
CAMERAS = ('1201-1', '1201-2')
SIDES = ('left', 'right')


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 verify(inputs, submission):
    source, meta = read(inputs), read(submission / 'submission.json')
    if meta['schema_version'] != SCHEMA:
        raise ValueError('expected v4 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')}
    expected.update(input_manifest_sha256=hashlib.sha256(inputs.read_bytes()).hexdigest(),
                    coordinate_frame='upright' 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) <= 200 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 {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 any(p.is_symlink() for p in submission.rglob('*')):
        raise ValueError('symlinks are not supported')
    for cid in ids:
        rows = [json.loads(line) for line in (submission / 'predictions' / (cid + '.jsonl')).read_text().splitlines()]
        if len(rows) != 150 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(150)):
            raise ValueError(cid + ': require exactly 150 unique frames')
        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 the upright original 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 hand.get(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) 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 = read(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', '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='upright' 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(150):
                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()