"""Author a FAITHFUL onnx-genai v1 inference_metadata.yaml for the real
gemma4-e2b target decoder, using the ACTUAL graph port names
(`past_key_values.N.key` / `present.N.key`), grouped into full/sliding state
services. Structure mirrors onnx-genai #1716 example 23; validated against that
PR's schema. Owner layers are read from the graph's per-port head_dim
(512=global/full, 256=local/sliding)."""
from __future__ import annotations
import json, sys
import onnx_ir as ir
import yaml

MODEL = "/datadisks/disk1/justinchu/inference-metadata-catalogue/gemma4/target/package/model.onnx"
SCHEMA = "/datadisks/disk1/justinchu/inference-metadata-catalogue/gemma4/ref-1716/schema.json"
OUT = "/datadisks/disk1/justinchu/inference-metadata-catalogue/gemma4/target/package/inference_metadata.yaml"

g = ir.load(MODEL).graph
ins = {v.name: v for v in g.inputs}
outs = {v.name: v for v in g.outputs}
# discover owner KV layers + their head_dim
owners = {}  # layer_idx -> head_dim
for name, v in ins.items():
    if name.startswith("past_key_values.") and name.endswith(".key"):
        idx = int(name.split(".")[1])
        hd = (v.shape[3].value if hasattr(v.shape[3], "value") else int(v.shape[3]))
        owners[idx] = hd
full = sorted(i for i, hd in owners.items() if hd == 512)
slide = sorted(i for i, hd in owners.items() if hd == 256)


def kv_contract(kind):
    hh, hd = (f"{kind}_kv_heads", f"{kind}_head_dim")
    return {"dtype": "float16", "rank": 4, "shape": ["batch", hh, "sequence", hd],
            "batch_layout": {"kind": "request_aligned", "axis": 0}}


def opaque_app(name):
    return {"role": {"kind": "opaque"}, "source": {"kind": "application", "name": name}, "required": True}


inputs = {
    "request.active": {"contract": {"dtype": "bool", "rank": 1, "shape": ["batch"], "batch_layout": {"kind": "request_aligned", "axis": 0}}, "role": {"kind": "opaque"}, "source": {"kind": "application", "name": "active"}, "required": True},
    "request.done": {"contract": {"dtype": "bool", "rank": 1, "shape": ["batch"], "batch_layout": {"kind": "request_aligned", "axis": 0}}, "role": {"kind": "opaque"}, "source": {"kind": "application", "name": "done"}, "required": True},
    "request.accepted_len": {"contract": {"dtype": "int64", "rank": 1, "shape": ["batch"], "batch_layout": {"kind": "request_aligned", "axis": 0}}, "role": {"kind": "opaque"}, "source": {"kind": "application", "name": "accepted_len"}, "required": True},
    "request.input_ids": {"contract": {"dtype": "int64", "rank": 2, "shape": ["batch", "sequence"], "batch_layout": {"kind": "request_aligned", "axis": 0}}, "role": {"kind": "runtime", "version": "1.0", "role": "prompt_tokens"}, "source": {"kind": "request"}, "required": True},
    "request.attention_mask": {"contract": {"dtype": "int64", "rank": 2, "shape": ["batch", "sequence"], "batch_layout": {"kind": "request_aligned", "axis": 0}}, **opaque_app("attention_mask")},
}
comp_in = {
    "input_ids": {"dtype": "int64", "rank": 2, "shape": ["batch", "sequence"], "batch_layout": {"kind": "request_aligned", "axis": 0}},
    "attention_mask": {"dtype": "int64", "rank": 2, "shape": ["batch", "sequence"], "batch_layout": {"kind": "request_aligned", "axis": 0}},
}
comp_out = {"logits": {"dtype": "float16", "rank": 3, "shape": ["batch", "sequence", "vocab"], "batch_layout": {"kind": "request_aligned", "axis": 0}}}
state = {}
step_in = {"input_ids": "request.input_ids", "attention_mask": "request.attention_mask"}
step_out = {"logits": "decoder.logits"}
groups = {"full_attention": {"kind": "full_attention", "sequence_axis": 2, "layout": "bnsh", "update": {"kind": "append"}, "reuse": {"prefix_reusable": True, "evictable_prefix": False}, "ports": {"decoder": {}}},
          "sliding_attention": {"kind": "sliding_attention", "sequence_axis": 2, "layout": "bnsh", "update": {"kind": "append"}, "reuse": {"prefix_reusable": True, "evictable_prefix": True}, "ports": {"decoder": {}}}}

for idx in sorted(owners):
    kind = "full" if owners[idx] == 512 else "sliding"
    grp = f"{kind}_attention"
    for role in ("key", "value"):
        pin = f"past_key_values.{idx}.{role}"
        pout = f"present.{idx}.{role}"
        inputs[f"request.{pin}"] = {"contract": kv_contract(kind), **opaque_app(pin)}
        comp_in[pin] = kv_contract(kind)
        comp_out[pout] = kv_contract(kind)
        state[pin] = {"contract": kv_contract(kind), "scope": "invocation", "initializer": f"request.{pin}",
                      "recurrence": {"kind": "invariant"}, "management": "runtime", "release_boundary": "invocation",
                      "service_group": grp}
        step_in[pin] = f"request.{pin}"
        step_out[pout] = f"decoder.{pin}"
        groups[grp]["ports"]["decoder"][pin] = {"input": pin, "output": pout, "role": role, "layer": idx}

doc = {
    "schema_version": "v1",
    "pipeline": {"workflow": {
        "manifest": {"capabilities": ["workflow_ssa", "typed_emit", "serving_service_contract"]},
        "inputs": inputs,
        "outputs": {"logits": {"contract": {"dtype": "float16", "rank": 3, "shape": ["batch", "sequence", "vocab"], "batch_layout": {"kind": "request_aligned", "axis": 0}}, "role": "tensor", "stage": "pre_adapter"}},
        "components": {"decoder": {"implementation": {"kind": "onnx", "artifact": "model.onnx"},
                                   "ports": {"inputs": comp_in, "outputs": comp_out, "roles": {"input_ids": "token_ids", "logits": "logits"}}}},
        "state": state,
        "steps": [{"kind": "invoke", "component": "decoder", "inputs": step_in, "outputs": step_out},
                  {"kind": "emit", "value": "decoder.logits", "output": "logits", "mode": "replace"}],
        "serving": {"active": "request.active", "done": "request.done", "accepted_len": "request.accepted_len",
                    "state_service": {"groups": groups}},
    }},
}

# validate
from jsonschema import Draft202012Validator
schema = json.load(open(SCHEMA))
errs = sorted(Draft202012Validator(schema).iter_errors(doc), key=lambda e: list(e.path))
if errs:
    print(f"INVALID ({len(errs)} errors):")
    for e in errs[:12]:
        print("  -", list(e.path), "->", e.message[:160])
    sys.exit(1)
header = ("# Faithful onnx-genai v1 inference_metadata for google/gemma-4-E2B-it text decoder.\n"
          "# Authored to the onnx-genai #1716 schema (example 23 shape) using the REAL exported\n"
          "# graph port names (past_key_values.N.key / present.N.key) and REAL owner layers:\n"
          f"#   full_attention owners (head_dim 512): layers {full}\n"
          f"#   sliding_attention owners (head_dim 256): layers {slide}\n"
          "# The 20 shared-KV layers borrow an owner's buffer inside the graph and expose no ports.\n"
          "# MoE is DISABLED in this checkpoint (enable_moe_block=false) -> dense MLP, not invented.\n"
          "# final_logit_softcapping / tie_word_embeddings are graph-internal. Validates against\n"
          "# onnx-genai PR #1716 schema/inference_metadata.schema.json.\n")
with open(OUT, "w") as f:
    f.write(header)
    yaml.safe_dump(doc, f, sort_keys=False, default_flow_style=None, width=1000)
print(f"VALID. full owners={full} sliding owners={slide}. Wrote {OUT}")