#!/usr/bin/env python3
# mergeSurvey.py - merge survey/*.json (phase 1 agent output) into annotation files and a knowledge base.
#
#   annotations/10_survey.json      unitLabels + routine header notes + text dataTypes
#   survey/variableProposals.json   every proposed variable name per address (for the naming arbiter)
#   docs/knowledgeBase.md           human readable routine / data / insight list for phase 2
import os, re, json, glob, sys

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
SURVEY = os.path.join(ROOT, "survey")
ANN = os.path.join(ROOT, "annotations")
DOCS = os.path.join(ROOT, "docs")

RESERVED = {"a", "x", "y", "s", "sp", "and", "or", "not", "xor", "mod", "div"}
MNEMONICS = {"adc","and","asl","bcc","bcs","beq","bit","bmi","bne","bpl","brk","bvc","bvs","clc","cld","cli","clv","cmp",
             "cpx","cpy","dec","dex","dey","eor","inc","inx","iny","jmp","jsr","lda","ldx","ldy","lsr","nop","ora","pha",
             "php","pla","plp","rol","ror","rti","rts","sbc","sec","sed","sei","sta","stx","sty","tax","tay","tsx","txa","txs","tya"}

sys.path.insert(0, os.path.join(ROOT, "tools"))
from buildAliases import UNIT_ALIASES


def canonicalUnit(u):
    return UNIT_ALIASES.get(u, u)


def unitTag(unit):
    base = unit.split("/")[-1]
    m = re.search(r"_(A|B|T\d\d|T1S\d+|T1S18_20|8000|0200|0400|C000|FBB8|9300|F000|8800)$", base)
    if base.startswith("ovl_6F00_"): return base[-1]
    if base.startswith("ovl_E000_"): return base[-3:]
    if base.startswith("ovl_EC00_") or base.startswith("ovl_EE00_"): return base[4:8] + base[-3:]
    if base.startswith("boot/") or unit.startswith("boot/"): return "Boot"
    if unit.startswith("drive/"): return "Drv"
    return ""

SPLIT_HINTS = ("low byte", "high byte", "lo byte", "hi byte", "low bytes", "high bytes",
               "lo bytes", "hi bytes", "split", "lsb table", "msb table")


def isSplitByteTable(name, description):
    """A 'lo'/'hi' half of a split address table is a byte table, not a table of words."""
    if re.search(r"(Lo|Hi|Low|High|Lsb|Msb)$", name):
        return True
    text = description.lower()
    return any(h in text for h in SPLIT_HINTS)


def validName(n):
    return bool(re.match(r"^[A-Za-z_][A-Za-z0-9_]*$", n)) and n.lower() not in RESERVED and n.lower() not in MNEMONICS

def main():
    files = sorted(glob.glob(os.path.join(SURVEY, "*.json")))
    files = [f for f in files if not f.endswith("variableProposals.json")]
    unitLabels = {}            # unit -> {addrHex: name}
    notes = {}                 # "addr|unit" -> note
    dataTypes = {}
    hwNames = set()
    try:
        sys.path.insert(0, os.path.join(ROOT, "tools"))
        from hwSymbols import c64Symbols, DRIVE
        hwNames = set(c64Symbols().values()) | set(DRIVE.values())
    except Exception:
        pass
    nameOwner = {}             # name -> (unit, addr) first claim
    varProposals = {}          # addr -> list of {name, meaning, scope, confidence, chunk}
    kb = {}                    # unit -> list of routine dicts
    kbData = {}
    insights = {}
    problems = []
    seed = json.load(open(os.path.join(ANN, "00_seed.json")))
    for a, n in seed.get("labels", {}).items():
        nameOwner[n] = ("seed", a.upper())
    for unit, m in seed.get("unitLabels", {}).items():
        for a, n in m.items():
            nameOwner[n] = (unit, a.upper())
    for z in seed.get("zp", {}).values():
        nameOwner[z[0] if isinstance(z, list) else z] = ("seed-zp", None)

    def claim(name, unit, addr):
        """return a globally unique variant of name for (unit, addr)"""
        if not validName(name) or name in hwNames:
            name = re.sub(r"[^A-Za-z0-9_]", "", name) or f"label{addr}"
            if name[0].isdigit(): name = "l" + name
            if name in hwNames or name.lower() in MNEMONICS or name.lower() in RESERVED: name += "Routine"
        owner = nameOwner.get(name)
        if owner is None or owner == (unit, addr):
            nameOwner[name] = (unit, addr)
            return name
        # same address claimed from another unit / the seed: same thing, keep the name
        if owner[1] == addr:
            return name
        tag = unitTag(unit)
        cand = name + tag if tag and not name.endswith(tag) else name + "_" + addr
        if nameOwner.get(cand) not in (None, (unit, addr)):
            cand = name + "_" + addr
        nameOwner[cand] = (unit, addr)
        return cand

    for f in files:
        try:
            j = json.load(open(f))
        except Exception as e:
            problems.append(f"{os.path.basename(f)}: invalid JSON ({e})")
            continue
        unit = canonicalUnit(j.get("unit") or "?")
        chunk = j.get("chunk") or os.path.basename(f)
        for r in j.get("routines", []):
            addr = str(r.get("addr", "")).upper().replace("$", "").zfill(4)
            if not re.match(r"^[0-9A-F]{4}$", addr):
                problems.append(f"{chunk}: bad routine addr {r.get('addr')}"); continue
            name = claim(str(r.get("name", f"sub_{addr}")), unit, addr)
            unitLabels.setdefault(unit, {})[addr] = name
            head = [f"{name} - {r.get('summary','').strip()}"]
            if r.get("inputs"):  head.append(f"In:  {r['inputs'].strip()}")
            if r.get("outputs"): head.append(f"Out: {r['outputs'].strip()}")
            conf = r.get("confidence", "")
            if conf and conf != "high": head.append(f"(confidence: {conf})")
            notes[f"{addr}|{unit}"] = {"routine": head, "unit": unit}
            kb.setdefault(unit, []).append({"addr": addr, "name": name, "summary": r.get("summary",""),
                                            "inputs": r.get("inputs",""), "outputs": r.get("outputs",""), "confidence": conf})
        for d in j.get("dataBlocks", []):
            addr = str(d.get("addr", "")).upper().replace("$", "").zfill(4)
            if not re.match(r"^[0-9A-F]{4}$", addr):
                problems.append(f"{chunk}: bad data addr {d.get('addr')}"); continue
            name = claim(str(d.get("name", f"data{addr}")), unit, addr)
            unitLabels.setdefault(unit, {})[addr] = name
            length = int(d.get("length", 0) or 0)
            t = d.get("type", "unknown")
            if t == "text" and 0 < length <= 256:
                dataTypes[addr] = ["text", length]
            elif t == "addrTable" and 0 < length <= 512 and length % 2 == 0 and not isSplitByteTable(name, d.get("description", "")):
                dataTypes[addr] = ["addr", length]
            desc = d.get("description", "")
            key = f"{addr}|{unit}"
            notes.setdefault(key, {"unit": unit})
            notes[key]["block"] = [f"{name}: {t}, {length} bytes. {desc}".strip()]
            kbData.setdefault(unit, []).append({"addr": addr, "name": name, "type": t, "length": length, "description": desc})
        for v in j.get("variables", []):
            addr = str(v.get("addr", "")).upper().replace("$", "")
            if not re.match(r"^[0-9A-F]{2}$|^[0-9A-F]{4}$", addr):
                continue
            varProposals.setdefault(addr, []).append({"name": v.get("name"), "meaning": v.get("meaning"), "scope": v.get("scope"),
                                                      "confidence": v.get("confidence"), "chunk": chunk, "unit": unit})
        for m in j.get("misclassified", []):
            problems.append(f"{chunk}: misclassified {m}")
        for s in j.get("insights", []):
            insights.setdefault(unit, []).append(f"[{chunk}] {s}")

    out = {"_comment": "generated by tools/mergeSurvey.py from survey/*.json - do not edit by hand",
           "unitLabels": unitLabels,
           "notes": {k.split("|")[0]: v for k, v in notes.items()} if False else {},
           "dataTypes": dataTypes}
    # notes keyed by address must stay unit scoped: build.py keys notes by address only, so emit one file per unit
    perUnitNotes = {}
    for k, v in notes.items():
        addr, unit = k.split("|")
        perUnitNotes.setdefault(unit, {})[addr] = v
    json.dump({"_comment": out["_comment"], "unitLabels": unitLabels, "dataTypes": dataTypes}, open(os.path.join(ANN, "10_survey_labels.json"), "w"), indent=1)
    for unit, m in perUnitNotes.items():
        fn = "11_survey_notes_" + unit.replace("/", "_") + ".json"
        json.dump({"_comment": out["_comment"], "notes": m}, open(os.path.join(ANN, fn), "w"), indent=1)
    json.dump(varProposals, open(os.path.join(SURVEY, "variableProposals.json"), "w"), indent=1)

    with open(os.path.join(DOCS, "knowledgeBase.md"), "w") as f:
        f.write("# Modem Wars - routine and data knowledge base (phase 1 survey, auto-generated)\n\n")
        f.write("Names are unit-scoped: overlay variants at the same address have different code.\n\n")
        for unit in sorted(kb):
            f.write(f"## {unit}\n\n")
            f.write("| addr | name | summary | in | out | conf |\n|---|---|---|---|---|---|\n")
            for r in sorted(kb[unit], key=lambda r: r["addr"]):
                row = [r["addr"], r["name"], r["summary"], r["inputs"], r["outputs"], r["confidence"]]
                f.write("| " + " | ".join(str(x).replace("|", "/").replace("\n", " ") for x in row) + " |\n")
            if unit in kbData:
                f.write("\nData:\n\n")
                for d in sorted(kbData[unit], key=lambda d: d["addr"]):
                    f.write(f"- `${d['addr']}` **{d['name']}** ({d['type']}, {d['length']} bytes): {d['description']}\n")
            if unit in insights:
                f.write("\nInsights:\n\n")
                for s in insights[unit]:
                    f.write(f"- {s}\n")
            f.write("\n")
    print(f"units: {len(unitLabels)}, labels: {sum(len(v) for v in unitLabels.values())}, notes: {len(notes)}, "
          f"variable addresses: {len(varProposals)}, problems: {len(problems)}")
    for p in problems[:40]:
        print("  -", p)

main()
