"""Versioned ONNX policies. Static admission checks, not a security sandbox."""
import io
import json
import stat
import zipfile
import numpy as np

ACTIONS = ['noop', 'left', 'right', 'up', 'down', 'drop']
MANIFEST = {
    'format': 'claw-policy-v1', 'model': 'policy.onnx',
    'input': {'name': 'frames', 'shape': [1, 4, 3, 120, 160], 'dtype': 'float32',
              'range': [0, 1], 'color': 'RGB', 'frame_order': 'oldest_first'},
    'output': {'name': 'action_logits', 'shape': [1, 6], 'dtype': 'float32'},
    'actions': ACTIONS, 'decision_hz': 10,
}
MAX_BYTES = 25 * 1024 * 1024
MAX_MEMORY_BYTES = 1024 * 1024
MIN_HOLD_MS, MAX_HOLD_MS = 1, 1000
V2 = 'claw-policy-v2'
ALLOWED_OPS = set('Conv Relu LeakyRelu Sigmoid Tanh Add Sub Mul Div MatMul Gemm Flatten Reshape Transpose Concat Slice Gather Constant Identity Cast ReduceMean ReduceSum ReduceMax Softmax MaxPool AveragePool GlobalAveragePool BatchNormalization Clip Abs Neg Squeeze Unsqueeze Shape Expand Where Greater Less Equal Pow Sqrt Resize Pad'.split())
# Fixed-shape recurrent cells and attention decompositions; no loops, custom
# kernels, external data or Python. Admission uses an isolated ARM64 server.
ALLOWED_OPS |= set('GRU LSTM RNN LayerNormalization Erf Exp Log Sin Cos Min Max And Or Not GreaterOrEqual LessOrEqual ArgMax GatherElements Split Tile ConstantOfShape CastLike'.split())


def manifest_v2(memory_elements=256):
    if type(memory_elements) is not int or not 1 <= memory_elements <= MAX_MEMORY_BYTES // 4:
        raise ValueError('Memory must contain 1..262144 float32 elements (1 MiB)')
    def tensor(name, shape): return dict(name=name, shape=shape, dtype='float32')
    return dict(format=V2, model='policy.onnx',
        inputs=[dict(tensor('image',[1,3,120,160]),range=[0,1],color='RGB'),
                tensor('previous_command',[1,7]),tensor('memory',[1,memory_elements])],
        outputs=[tensor('action_logits',[1,6]),tensor('hold_ms',[1,1]),
                 tensor('next_memory',[1,memory_elements])],
        actions=ACTIONS, camera_hz=10, memory_elements=memory_elements,
        direction_hold_ms=dict(min=MIN_HOLD_MS,max=MAX_HOLD_MS), drop_hold_ms=250,
        previous_command_encoding='six_action_one_hot_then_executed_hold_ms;zero_at_start',
        scheduling='complete_hold_release_then_fresh_image')


def expected_manifest(manifest):
    if manifest.get('format') == MANIFEST['format']: expected = MANIFEST
    elif manifest.get('format') == V2: expected = manifest_v2(manifest.get('memory_elements'))
    else: raise ValueError('Unsupported policy format')
    if any(manifest.get(k) != v for k,v in expected.items()):
        raise ValueError('Manifest does not match '+expected['format'])
    return expected


def command_duration(action, requested_ms, remaining_ms=None):
    """Host authority: finite, integral milliseconds, clipped to the profile."""
    if type(action) is not int or not 0 <= action < len(ACTIONS) or not np.isfinite(requested_ms):
        raise ValueError('Invalid action/duration')
    if action == 0: return 0
    if action == 5: return 250
    duration = int(np.floor(np.clip(requested_ms, MIN_HOLD_MS, MAX_HOLD_MS)+.5))
    if remaining_ms is not None:
        if not np.isfinite(remaining_ms): raise ValueError('Invalid remaining budget')
        duration = min(duration, max(0, int(np.floor(remaining_ms))))
    return duration


def command_feedback(action=None, held_ms=0.):
    result=np.zeros((1,7),np.float32)
    if action is not None:
        if type(action) is not int or not 0 <= action < 6 or not np.isfinite(held_ms) or held_ms < 0:
            raise ValueError('Invalid command feedback')
        result[0,action]=1; result[0,6]=held_ms
    return result


def validate(path):
    import onnx
    from pathlib import Path
    if Path(path).stat().st_size > MAX_BYTES:
        raise ValueError('Submission exceeds 25 MiB')
    with zipfile.ZipFile(path) as z:
        infos = z.infolist()
        if len(infos) != 2 or {i.filename for i in infos} != {'manifest.json', 'policy.onnx'}:
            raise ValueError('Exactly manifest.json and policy.onnx required')
        if sum(i.file_size for i in infos) > MAX_BYTES:
            raise ValueError('Uncompressed submission exceeds 25 MiB')
        if any(stat.S_ISLNK(i.external_attr >> 16) or i.flag_bits & 1 for i in infos):
            raise ValueError('Links and encrypted members prohibited')
        manifest = json.loads(z.read('manifest.json'))
        expected = expected_manifest(manifest)
        blob = z.read('policy.onnx')
    graph = onnx.load_model_from_string(blob)
    if graph.functions or len(graph.graph.node) > 2000:
        raise ValueError('Local functions or excessive graph size')
    if any(n.domain not in ('', 'ai.onnx') or n.op_type not in ALLOWED_OPS for n in graph.graph.node):
        raise ValueError('Unsupported graph operator')
    if any(a.type in (onnx.AttributeProto.GRAPH, onnx.AttributeProto.GRAPHS) for n in graph.graph.node for a in n.attribute):
        raise ValueError('Nested graphs prohibited')
    tensors = list(graph.graph.initializer)
    for n in graph.graph.node:
        for a in n.attribute:
            if a.type == onnx.AttributeProto.TENSOR:
                tensors.append(a.t)
            tensors.extend(a.tensors)
    if graph.graph.sparse_initializer:
        raise ValueError('Sparse initializers prohibited')
    for tensor in tensors:
        if tensor.external_data or tensor.data_location == onnx.TensorProto.EXTERNAL:
            raise ValueError('External tensors prohibited')
        if not np.isfinite(onnx.numpy_helper.to_array(tensor)).all():
            raise ValueError('Nonfinite model weights')
    onnx.checker.check_model(graph, full_check=True)
    inputs=expected.get('inputs',[expected.get('input')]); outputs=expected.get('outputs',[expected.get('output')])
    for actual, specifications in ((graph.graph.input, inputs), (graph.graph.output, outputs)):
        if len(actual) != len(specifications): raise ValueError('Incorrect number of graph inputs/outputs')
        for value, spec in zip(actual,specifications):
            t=value.type.tensor_type
            if value.name != spec['name'] or t.elem_type != onnx.TensorProto.FLOAT or [d.dim_value for d in t.shape.dim] != spec['shape']:
                raise ValueError('Incorrect input/output tensor contract')
    inferred = onnx.shape_inference.infer_shapes(graph, strict_mode=True)
    activation_elements = 0
    values={v.name:v for v in list(inferred.graph.value_info)+list(inferred.graph.input)+list(inferred.graph.output)}
    produced={name for node in inferred.graph.node for name in node.output if name}
    if not produced <= values.keys(): raise ValueError('Missing inferred activation shapes')
    for value in (values[name] for name in produced):
        t = value.type.tensor_type
        if not t.HasField('shape') or any(d.dim_param or not d.HasField('dim_value') or d.dim_value < 0 for d in t.shape.dim):
            raise ValueError('Unresolved activation shape')
        activation_elements += int(np.prod([d.dim_value for d in t.shape.dim], dtype=object))
    if activation_elements > 64_000_000:
        raise ValueError('Activation budget exceeded')
    return manifest, blob


def package(onnx_path, output, *, version=1, memory_elements=256, **metadata):
    if version not in (1,2): raise ValueError('Unsupported policy version')
    specification=MANIFEST if version == 1 else manifest_v2(memory_elements)
    with zipfile.ZipFile(output, 'w', compression=zipfile.ZIP_DEFLATED) as z:
        z.writestr('manifest.json', json.dumps({**metadata, **specification}, indent=2)+'\n')
        z.write(onnx_path, 'policy.onnx')
    validate(output)


def make_session(blob):
    import onnxruntime as ort
    # The macOS telemetry worker crashed after a completed training export.
    # Disable it before creating sessions; it is unrelated to policy inference.
    ort.disable_telemetry_events()
    options = ort.SessionOptions()
    options.intra_op_num_threads = 1
    options.inter_op_num_threads = 1
    return ort.InferenceSession(blob, sess_options=options, providers=['CPUExecutionProvider'])


def infer(session, frames):
    logits = session.run(['action_logits'], {'frames': frames})[0]
    if logits.shape != (1, 6) or logits.dtype != np.float32 or not np.isfinite(logits).all():
        raise ValueError('Invalid policy output')
    return logits[0]


class PolicyRuntime:
    """Trusted wrapper; weights fixed, opaque memory owned by this episode."""
    def __init__(self, path):
        self.manifest,blob=validate(path); self.session=make_session(blob)
        self.format=self.manifest['format']; self.reset()

    def reset(self):
        self.memory=np.zeros((1,self.manifest.get('memory_elements',1)),np.float32)

    @property
    def image_shape(self):
        return (1,3,120,160) if self.format == V2 else (1,4,3,120,160)

    def step(self, image, previous_command=None):
        image=np.asarray(image)
        if image.shape != self.image_shape or image.dtype != np.float32 or not np.isfinite(image).all() or image.min()<0 or image.max()>1:
            raise ValueError('Invalid image tensor')
        if self.format != V2: return infer(self.session,image),75.
        previous_command=command_feedback() if previous_command is None else np.asarray(previous_command)
        if previous_command.shape != (1,7) or previous_command.dtype != np.float32 or not np.isfinite(previous_command).all():
            raise ValueError('Invalid previous command tensor')
        logits,hold,memory=self.session.run(['action_logits','hold_ms','next_memory'],
            dict(image=image,previous_command=previous_command,memory=self.memory))
        for value,shape in ((logits,(1,6)),(hold,(1,1)),(memory,self.memory.shape)):
            if value.shape != shape or value.dtype != np.float32 or not np.isfinite(value).all():
                raise ValueError('Invalid policy output or memory')
        self.memory=memory.copy()
        return logits[0],float(hold[0,0])

    def preflight(self, steps=32, limit_ms=75.):
        """Bounded representative admission check; every live step is timed too."""
        import time
        rng=np.random.default_rng(1701); timings=[]
        self.reset()
        self.step(np.zeros(self.image_shape,np.float32),command_feedback())
        for i in range(steps):
            image=(np.zeros(self.image_shape,np.float32) if i%3==0 else
                   np.ones(self.image_shape,np.float32) if i%3==1 else
                   rng.random(self.image_shape,dtype=np.float32))
            start=time.monotonic_ns()
            scores,hold=self.step(image,command_feedback(i%6,75 if i%6 in range(1,5) else 0))
            elapsed=(time.monotonic_ns()-start)/1e6; timings.append(elapsed)
            if elapsed>limit_ms: raise ValueError('Model exceeds admission inference budget')
        self.reset()
        return dict(steps=steps,max_ms=max(timings),p95_ms=float(np.quantile(timings,.95)),
                    limit_ms=limit_ms,memory_bytes=self.memory.nbytes if self.format==V2 else 0)


if __name__ == '__main__':
    import sys
    runtime=PolicyRuntime(sys.argv[1])
    print(json.dumps({'manifest':runtime.manifest,'preflight':runtime.preflight()}))
