"""Public C API real-sample regression against a saved geometry comparison run."""
import argparse
import ctypes as C
import hashlib
import json
import os
from pathlib import Path

import cv2
import numpy as np

class Image(C.Structure):
    _fields_ = [('struct_size', C.c_uint32), ('version', C.c_uint32), ('data', C.c_void_p),
                ('width', C.c_int32), ('height', C.c_int32), ('stride', C.c_int32),
                ('format', C.c_int32), ('reserved', C.c_uint32 * 8)]

class Options(C.Structure):
    _fields_ = [('struct_size', C.c_uint32), ('version', C.c_uint32), ('mode', C.c_int32), ('reserved', C.c_uint32 * 13)]

class Error(C.Structure):
    _fields_ = [('struct_size', C.c_uint32), ('version', C.c_uint32), ('status', C.c_int32), ('message', C.c_char * 512), ('reserved', C.c_uint32 * 8)]

class Info(C.Structure):
    _fields_ = [('struct_size', C.c_uint32), ('version', C.c_uint32)] + [(n, C.c_uint32) for n in
        ['abi', 'task', 'capabilities', 'formats', 'modes', 'max_results', 'max_keypoints']] + [('reserved', C.c_uint32 * 8)]

class Geometry(C.Structure):
    _fields_ = [('struct_size', C.c_uint32), ('version', C.c_uint32), ('type', C.c_int32)] + [(n, C.c_float) for n in
        ['x', 'y', 'width', 'height', 'angle', 'radius']] + [('fit_status', C.c_uint32), ('reserved', C.c_uint32 * 7)]

class Mask(C.Structure):
    _fields_ = [('struct_size', C.c_uint32), ('version', C.c_uint32), ('data', C.c_void_p)] + [(n, C.c_int32) for n in
        ['width', 'height', 'stride', 'x', 'y']] + [('reserved', C.c_uint32 * 8)]

def versioned(cls):
    value = cls(); value.struct_size = C.sizeof(cls); value.version = 1
    return value

def main():
    p = argparse.ArgumentParser()
    p.add_argument('--install', type=Path, required=True)
    p.add_argument('--evidence', type=Path, required=True)
    p.add_argument('--reference', type=Path, action='append', required=True,
                   help='Prior run directory containing raw.json and provenance.json; tests its first image')
    args = p.parse_args(); args.evidence.mkdir(parents=True, exist_ok=True)
    libdir = args.install.resolve() / 'lib'
    cookie = os.add_dll_directory(str(libdir)) if os.name == 'nt' else None
    lib = C.CDLL(str(libdir / ('yolo_runtime.dll' if os.name == 'nt' else 'libyolo_runtime.so')))
    signatures = {
        'yolo_runtime_create': ([C.c_char_p, C.POINTER(C.c_void_p), C.POINTER(Error)], C.c_int32),
        'yolo_runtime_get_info': ([C.c_void_p, C.POINTER(Info)], C.c_int32),
        'yolo_runtime_infer_image': ([C.c_void_p, C.POINTER(Image), C.POINTER(Options), C.POINTER(C.c_void_p), C.POINTER(Error)], C.c_int32),
        'yolo_result_set_get_count': ([C.c_void_p, C.POINTER(C.c_uint32)], C.c_int32),
        'yolo_result_set_get_seg_geometry': ([C.c_void_p, C.c_uint32, C.POINTER(Geometry)], C.c_int32),
        'yolo_result_set_get_seg_mask': ([C.c_void_p, C.c_uint32, C.POINTER(Mask)], C.c_int32),
        'yolo_result_set_release': ([C.c_void_p], None), 'yolo_runtime_destroy': ([C.c_void_p], None)}
    for name, (types, result) in signatures.items():
        f = getattr(lib, name); f.argtypes = types; f.restype = result
    assert C.sizeof(Geometry) == 68 and Geometry.radius.offset == 32 and Geometry.fit_status.offset == 36
    reports = []
    for ri, reference in enumerate(args.reference):
        provenance = json.loads((reference / 'provenance.json').read_text(encoding='utf-8'))
        raw = json.loads((reference / 'raw.json').read_text(encoding='utf-8'))
        first = raw['images'][0]
        model = Path(provenance['model']); image_path = Path(provenance['input']) / first['image']
        assert hashlib.sha256(model.read_bytes()).hexdigest() == provenance['model_sha256']
        # Labels do not affect segmentation geometry; preserve the actual model metadata.
        import onnx
        import ast
        metadata = {v.key: v.value for v in onnx.load(str(model), load_external_data=False).metadata_props}
        labels = list(ast.literal_eval(metadata['names']).values())
        config = args.evidence / f'config-{ri}.yaml'
        config.write_text(f'''schema_version: 1
name: four_geometry_test
version: 1
task: seg
model: {json.dumps(str(model), ensure_ascii=False)}
target_platform: rk3588
model_family: yolo8
output_contract: task_specific
input:
  width: {raw['imgsz']}
  height: {raw['imgsz']}
  format: rgb
  layout: nchw
  normalization:
    mean: [0, 0, 0]
    std: [255, 255, 255]
postprocess:
  adapter: ultralytics_yolo8_seg
  confidence_threshold: {raw['conf']}
  nms_threshold: {raw['iou']}
  class_agnostic_nms: false
labels: {json.dumps(labels)}
''', encoding='utf-8')
        runtime = C.c_void_p(); error = versioned(Error)
        assert lib.yolo_runtime_create(str(config.resolve()).encode('utf-8'), C.byref(runtime), C.byref(error)) == 0, bytes(error.message)
        try:
            info = versioned(Info); assert lib.yolo_runtime_get_info(runtime, C.byref(info)) == 0
            assert info.abi == 5 and info.modes == 30 and info.capabilities & 30 == 30
            image = cv2.imdecode(np.fromfile(str(image_path), np.uint8), cv2.IMREAD_COLOR)
            image = np.ascontiguousarray(image)
            view = versioned(Image); view.data = image.ctypes.data; view.width = image.shape[1]; view.height = image.shape[0]
            view.stride = image.strides[0]; view.format = 2
            maxima = {mode: 0.0 for mode in range(1, 5)}
            for mode in [1, 2, 3, 4, 0, 5]:
                options = versioned(Options); options.mode = mode; results = C.c_void_p()
                status = lib.yolo_runtime_infer_image(runtime, C.byref(view), C.byref(options), C.byref(results), C.byref(error))
                if mode in (0, 5):
                    assert status == -1 and not results.value
                    continue
                assert status == 0, bytes(error.message)
                try:
                    count = C.c_uint32(); assert lib.yolo_result_set_get_count(results, C.byref(count)) == 0
                    assert count.value == len(first['targets'])
                    for i, expected in enumerate(first['targets']):
                        g = versioned(Geometry); m = versioned(Mask)
                        assert lib.yolo_result_set_get_seg_geometry(results, i, C.byref(g)) == 0
                        assert lib.yolo_result_set_get_seg_mask(results, i, C.byref(m)) == 0
                        expected_status = (1 if expected['fitted'] else 2) if mode == 4 else 0
                        assert g.type == mode and g.fit_status == expected_status
                        local = np.ctypeslib.as_array((C.c_uint8 * (m.stride*m.height)).from_address(m.data)).reshape(m.height, m.stride)[:, :m.width]
                        full = np.zeros(image.shape[:2], np.uint8); full[m.y:m.y+m.height, m.x:m.x+m.width] = local
                        assert hashlib.sha256(full.tobytes()).hexdigest() == expected['mask_sha256']
                        if mode == 4: want = [expected['fit_x'], expected['fit_y'], expected['fit_r']]; got = [g.x, g.y, g.radius]
                        elif mode == 3: want = [expected['mean_x'], expected['mean_y'], expected['mean_r']]; got = [g.x, g.y, g.radius]
                        elif mode == 1: want = [expected['mean_x'], expected['mean_y'], expected['rect_w'], expected['rect_h']]; got = [g.x, g.y, g.width, g.height]
                        else:
                            yy, xx = np.nonzero(full)
                            center, radius = cv2.minEnclosingCircle(np.column_stack([xx, yy]).astype(np.float32))
                            want = [*center, radius]; got = [g.x, g.y, g.radius]
                        diff = max(abs(a-b) for a,b in zip(got,want)); assert diff < .002, (mode, i, diff)
                        maxima[mode] = max(maxima[mode], diff)
                finally:
                    lib.yolo_result_set_release(results)
            reports.append(dict(reference=str(reference), image=first['image'], targets=len(first['targets']), max_error_px=maxima,
                                masks_unchanged=True, modes=[1,2,3,4], invalid_modes_rejected=[0,5]))
        finally:
            lib.yolo_runtime_destroy(runtime)
    (args.evidence / 'four-mode-api.json').write_text(json.dumps(reports, ensure_ascii=False, indent=2), encoding='utf-8')
    print('FOUR_GEOMETRY_PUBLIC_API_PASS', json.dumps(reports, ensure_ascii=False))

if __name__ == '__main__':
    main()