import json, sys

src, dst = sys.argv[1], sys.argv[2]
with open(src, "r", encoding="utf-8") as f:
    d = json.load(f)

# Racine / scalars
d.setdefault("version", "0.1.1")
d.setdefault("meta", {"source_lang":"AXIR"})
d.setdefault("types", {}); d["types"].setdefault("scalars", {})
sc = d["types"]["scalars"]

def ensure_scalar(name, dtype, val):
    sc[name] = {"dtype": dtype, "value": val}

N = int(sc.get("N", {}).get("value", 65536))
alpha = float(sc.get("alpha", {}).get("value", 1.0))
ensure_scalar("N", "i32", N)
ensure_scalar("alpha", "float32", alpha)

# Buffers host
d.setdefault("buffers", {})
buf = d["buffers"]
buf.setdefault("hA", {"role":"host","dtype":"float32","size":"N"})
buf.setdefault("hB", {"role":"host","dtype":"float32","size":"N"})
buf.setdefault("hC", {"role":"host","dtype":"float32","size":"N"})

ops = d.setdefault("ops", [])

def has_op(t, **kv):
    for op in ops:
        if op.get("op")==t and all(op.get(k)==v for k,v in kv.items()):
            return True
    return False

# Préfixes nécessaires
prepend = []
if not any(op.get("op")=="HostMake" and op.get("name")=="hA" for op in ops):
    prepend.append({"op":"HostMake","name":"hA","dtype":"float32","shape":["N"],"fill":"linspace","start":1.0,"step":1.0})
if not any(op.get("op")=="HostMake" and op.get("name")=="hB" for op in ops):
    prepend.append({"op":"HostMake","name":"hB","dtype":"float32","shape":["N"],"fill":"ones"})
if not any(op.get("op")=="HostMake" and op.get("name")=="hC" for op in ops):
    prepend.append({"op":"HostMake","name":"hC","dtype":"float32","shape":["N"],"fill":"zeros"})

if not has_op("DeviceMallocLike", dst="&dA"):
    prepend.append({"op":"DeviceMallocLike","dst":"&dA","like":"hA"})
if not has_op("DeviceMallocLike", dst="&dB"):
    prepend.append({"op":"DeviceMallocLike","dst":"&dB","like":"hB"})
if not has_op("DeviceMallocLike", dst="&dC"):
    prepend.append({"op":"DeviceMallocLike","dst":"&dC","like":"hC"})

# Recompose ops
ops = prepend + ops

# Memcpy H2D
insert_at = len(prepend)
if not has_op("Memcpy", dst="dA", kind="H2D"):
    ops.insert(insert_at, {"op":"Memcpy","dst":"dA","src":"hA","kind":"H2D"}); insert_at += 1
if not has_op("Memcpy", dst="dB", kind="H2D"):
    ops.insert(insert_at, {"op":"Memcpy","dst":"dB","src":"hB","kind":"H2D"}); insert_at += 1

# KernelLaunch saxpy si absent
if not any(op.get("op")=="KernelLaunch" and op.get("kernel")=="saxpy" for op in ops):
    ops.append({
        "op":"KernelLaunch",
        "kernel":"saxpy",
        "grid":["((N + 255) / 256)", 1, 1],
        "block":[256, 1, 1],
        "args":["dA","dB","dC","alpha","N"]
    })

# Memcpy D2H
if not has_op("Memcpy", dst="hC", kind="D2H"):
    ops.append({"op":"Memcpy","dst":"hC","src":"dC","kind":"D2H"})

d["ops"] = ops
with open(dst, "w", encoding="utf-8") as f:
    json.dump(d, f, indent=2)
print("fixed ->", dst)
