#!/usr/bin/env python3
"""Check the normalized example render format; not a visual-quality audit."""
from __future__ import annotations
import argparse
import ast
import math
import re
import struct
import zlib
from pathlib import Path
from _common import read_json, selected_ids, write_csv


def png_size(path: Path) -> tuple[int, int] | None:
    try:
        with path.open('rb') as f:
            head = f.read(33)
        if len(head) != 33 or head[:8] != b'\x89PNG\r\n\x1a\n' or head[8:16] != b'\x00\x00\x00\rIHDR':
            return None
        if zlib.crc32(head[12:29]) & 0xffffffff != struct.unpack('>I', head[29:33])[0]:
            return None
        width, height = struct.unpack('>II', head[16:24])
        return (height, width) if width > 0 and height > 0 else None
    except (OSError, ValueError, struct.error):
        return None


def npy_size(path: Path) -> tuple[int, int] | None:
    try:
        with path.open('rb') as f:
            if f.read(6) != b'\x93NUMPY':
                return None
            version = f.read(2)
            if version not in (b'\x01\x00', b'\x02\x00', b'\x03\x00'):
                return None
            size_fmt = '<H' if version == b'\x01\x00' else '<I'
            length = struct.unpack(size_fmt, f.read(struct.calcsize(size_fmt)))[0]
            if length > 1024 * 1024:
                return None
            header = ast.literal_eval(f.read(length).decode('utf-8' if version == b'\x03\x00' else 'latin1').strip())
            if not isinstance(header, dict):
                return None
            shape = header.get('shape')
            dtype = re.fullmatch(r'[<>=|]f([48])', str(header.get('descr', '')))
            if not dtype or not isinstance(shape, tuple) or len(shape) != 2 or not all(type(v) is int and v > 0 for v in shape):
                return None
            if not isinstance(header.get('fortran_order'), bool):
                return None
            expected = shape[0] * shape[1] * int(dtype.group(1))
            if path.stat().st_size - f.tell() != expected:
                return None
            return shape
    except (OSError, ValueError, SyntaxError, struct.error, UnicodeError):
        return None


def has_magic(path: Path, magic: bytes) -> bool:
    try:
        with path.open('rb') as f:
            return f.read(len(magic)) == magic
    except OSError:
        return False


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--render-dir', type=Path, required=True)
    parser.add_argument('--metadata', type=Path, required=True)
    parser.add_argument('--output', type=Path, required=True)
    parser.add_argument('--limit', type=int, default=50)
    args = parser.parse_args()
    if args.limit < 1:
        parser.error('--limit must be positive')
    doc = read_json(args.metadata)
    ids = selected_ids(doc)
    if not ids:
        raise SystemExit('Cannot audit an empty selection')
    rows = []
    for item in doc['selected_viewpoints'][:args.limit]:
        cid = item['candidate_id']
        if cid in ('.', '..') or any(v in cid for v in ('/', '\\')) or type(item.get('rank')) is not int or item['rank'] < 0:
            raise SystemExit('Candidate ID and rank must form a safe normalized render filename')
        rank = item['rank']
        stem = f'{rank:03d}_{cid}'
        rgb = args.render_dir / f'{stem}_rgb.png'
        depth = args.render_dir / f'{stem}_depth.npy'
        pose = args.render_dir / f'{stem}_pose.json'
        rgb_shape, depth_shape = png_size(rgb), npy_size(depth)
        pose_ok = False
        try:
            value = read_json(pose)
            position = value.get('position') if isinstance(value, dict) else None
            pose_ok = isinstance(position, list) and len(position) == 3 and all(
                type(v) in (int, float) and math.isfinite(v) for v in position
            ) and value.get('candidate_id') == cid
        except (OSError, ValueError):
            pass
        checks = {
            'rgb_exists': rgb.is_file(), 'rgb_png_magic': has_magic(rgb, b'\x89PNG\r\n\x1a\n'),
            'rgb_header_valid': rgb_shape is not None,
            'depth_exists': depth.is_file(), 'depth_npy_magic': has_magic(depth, b'\x93NUMPY'),
            'depth_array_valid': depth_shape is not None,
            'rgb_depth_shape_match': rgb_shape is not None and rgb_shape == depth_shape,
            'pose_exists': pose.is_file(), 'pose_json_valid': pose_ok,
        }
        rows.append({'candidate_id': cid, 'rank': rank, **checks, 'passed': all(checks.values())})
    write_csv(args.output, rows)
    failed = sum(not row['passed'] for row in rows)
    print(f'Wrote {args.output}: {len(rows) - failed}/{len(rows)} file/schema checks passed')
    if failed:
        raise SystemExit(1)


if __name__ == '__main__':
    main()