#!/usr/bin/env python3
"""Extract a *INIT data table to JSON. Auto-detects the table's shape.

*INIT scripts populate parallel global arrays with static game data. Three shapes seen:

  name   — records keyed by a name string. Each record: set-string(name), static field writes,
           set-string(desc). Arrays indexed by record id in lockstep (+1/record).
           (SKINIT skills, ITINIT items, EBINIT units)
  numeric— column table with NO names: mov/copy-to-global into parallel int arrays, keyed
           by an incrementing index column. (CGINIT gallery)
  footer — copy-local-array (op 0x64) bulk-loads length-prefixed arrays from the file
           footer into per-record global arrays. The data lives in the footer. (MPINIT maps)

Records are {id, name?, desc?, fields:{"0x<col_base>": value}} or, for footer tables,
{id, global_addr, footer_off, values:[...]}. Column addresses are raw engine globals;
naming them (attack, cost, …) needs the engine global-var map — later work.

Usage: py -3.11 -X utf8 tools/extract_init.py <TABLE> [OUTNAME] [--mode name|numeric|footer]
"""
from __future__ import annotations
import json
import sys
from functools import cache
from pathlib import Path

HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE))
import paths
import extract_message_table
import sys4load

SET_STRING = 0x192
MOV = 0x55
SUB = 0x51
COPY_TO_GLOBAL = 0x6C
COPY_LOCAL_ARRAY = 0x64
T_GLOBAL_INT = 3
T_GLOBAL_STRING = 5
T_IMM = 0

MESSAGE_TABLES = {
    "ITINIT": "ITMES",
    "SKINIT": "SKMES",
}


def resolve(name: str) -> Path:
    for cand in (paths.GAME_DIR / f"{name}.BIN", paths.DATA1 / f"{name}.BIN"):
        if cand.exists():
            return cand
    raise SystemExit(f"not found: {name}.BIN")


def _val(arg):
    """Render an operand as an int (immediate) or a {type,value} ref."""
    t, v = arg
    return v if t == T_IMM else {"type": f"0x{t:x}", "value": f"0x{v:x}"}


def _static_global_write(ins):
    """Return (destination, value) for statically evaluable global-int writes.

    The shipped name-mode INIT scripts encode positive values with `mov` and
    negative values with `sub destination, 0, magnitude`.  Ignoring the latter
    silently drops costs and penalties from the extracted schema.
    """
    if not ins.args or ins.args[0][0] != T_GLOBAL_INT:
        return None
    if ins.opcode == MOV and len(ins.args) >= 2:
        return ins.args[0][1], _val(ins.args[1])
    if (ins.opcode == SUB and len(ins.args) >= 3
            and ins.args[1][0] == T_IMM and ins.args[2][0] == T_IMM):
        return ins.args[0][1], ins.args[1][1] - ins.args[2][1]
    return None


def read_footer_array(scr, off):
    """Read a length-prefixed Data_Array at dword `off`: [length][v0..v_{length-1}]."""
    dw = scr.dwords
    if not (0 <= off < scr.nbody):
        return None
    length = dw[off]
    if length > scr.nbody or off + 1 + length > scr.nbody:
        return None
    return list(dw[off + 1: off + 1 + length])


def detect_mode(scr):
    ops = [ins.opcode for ins in scr.instructions]
    has_str = any(ins.opcode == SET_STRING and ins.args and ins.args[0][0] == T_GLOBAL_STRING
                  for ins in scr.instructions)
    if has_str:
        return "name"
    n_footer = ops.count(COPY_LOCAL_ARRAY)
    n_int = ops.count(MOV) + ops.count(COPY_TO_GLOBAL)
    return "footer" if n_footer >= max(4, n_int) else "numeric"


def _infer_record_span(string_addrs):
    """Infer the reserved width of one parallel string-array column.

    The shipped INIT tables reserve a fixed number of ids per column (300 for
    SKINIT and 1000 for ITINIT/EBINIT).  A populated record commonly writes its
    name and then its description, so that column stride is the dominant large
    positive delta between consecutive string destinations.
    """
    counts = {}
    for left, right in zip(string_addrs, string_addrs[1:]):
        delta = right - left
        if delta >= 32:
            counts[delta] = counts.get(delta, 0) + 1
    if not counts:
        raise ValueError("cannot infer name-table record span")
    return max(counts, key=lambda delta: (counts[delta], delta))


@cache
def _known_record_tables():
    """Return corpus-observed (base, stride) pairs used by lookup-array-2d.

    INIT scripts often populate linked row-major tables while defining an
    entity.  Treating every such write as `destination - entity_id` invents a
    different one-off parallel column for every row.  Consumer bytecode gives
    us the unambiguous table base and stride instead.
    """
    tables = set()
    for path in paths.scripts().values():
        try:
            script = sys4load.load(path)
        except sys4load.Sys4Error:
            continue
        for ins in script.instructions:
            if (sys4load.display_label(ins.opcode) == "lookup-array-2d"
                    and len(ins.args) >= 5
                    and ins.args[1][0] in (T_GLOBAL_INT, 6)
                    and ins.args[3][0] == T_IMM
                    and ins.args[3][1] > 0):
                tables.add((ins.args[1][1], ins.args[3][1]))
    return tuple(sorted(tables))


def _record_table_cell(destination, record_id):
    matches = []
    for base, stride in _known_record_tables():
        column = destination - (base + record_id * stride)
        if 0 <= column < stride:
            matches.append((base, stride, column))
    if len(matches) > 1:
        raise ValueError(
            f"ambiguous record-table destination 0x{destination:x} for id {record_id}: {matches}"
        )
    return matches[0] if matches else None


def extract_name(scr):
    string_addrs = [
        ins.args[0][1]
        for ins in scr.instructions
        if ins.opcode == SET_STRING and ins.args and ins.args[0][0] == T_GLOBAL_STRING
    ]
    if not string_addrs:
        return [], {}
    name_write_base = string_addrs[0]
    # AGE's shipped entity ids are one-based.  Array lookups use the cell just
    # before the first populated destination as their base, then add the id.
    first_record_id = 1
    name_base = name_write_base - first_record_id
    record_span = _infer_record_span(string_addrs)
    records, cur, desc_slot, desc_bases = [], None, 0, {}
    for ins in scr.instructions:
        if ins.opcode == SET_STRING and ins.args and ins.args[0][0] == T_GLOBAL_STRING:
            addr = ins.args[0][1]
            txt = scr.strings.get(ins.args[1][1], (None,))[0] if len(ins.args) > 1 else None
            # Names occupy column zero.  Do not use an address decrease as the
            # boundary: ITINIT begins with 101 consecutive name-only records,
            # which the old heuristic collapsed into item zero.
            if name_write_base <= addr < name_write_base + record_span:
                cur = {"id": addr - name_base, "name": txt, "fields": {}, "record_fields": {}}
                records.append(cur); desc_slot = 0
            elif cur is not None:
                key = "desc" if desc_slot == 0 else f"desc{desc_slot}"
                cur[key] = txt; desc_bases.setdefault(key, addr - cur["id"]); desc_slot += 1
        elif cur is not None and (write := _static_global_write(ins)) is not None:
            destination, value = write
            cell = _record_table_cell(destination, cur["id"])
            if cell is None:
                cur["fields"][f"0x{destination - cur['id']:x}"] = value
            else:
                base, stride, column = cell
                cur["record_fields"][f"0x{base:x}/{stride}/{column}"] = value
    for record in records:
        if not record["record_fields"]:
            del record["record_fields"]
    record_columns = sorted(
        {key for record in records for key in record.get("record_fields", {})},
        key=lambda key: tuple(int(part, 0) for part in key.split("/")),
    )
    return records, {"name_array_base": f"0x{name_base:x}",
                     "name_write_base": f"0x{name_write_base:x}",
                     "first_record_id": first_record_id,
                     "record_span": record_span,
                     "record_field_columns": record_columns,
                     "desc_array_bases": {k: f"0x{v:x}" for k, v in sorted(desc_bases.items())}}


def _int_writes(scr):
    """Ordered (addr, value_arg) for global-int mov / copy-to-global."""
    out = []
    for ins in scr.instructions:
        if ins.opcode in (MOV, COPY_TO_GLOBAL) and ins.args and ins.args[0][0] == T_GLOBAL_INT:
            out.append((ins.args[0][1], ins.args[1]))
    return out


def _longest_stride1_column(addrs):
    """Pick the primary index array: the stride-1 arithmetic run covering the most records."""
    seen = set(addrs)
    best_base, best_len = None, 0
    for a in sorted(seen):
        if a - 1 in seen:
            continue                       # only start at a run's base
        n = 0
        while a + n in seen:
            n += 1
        if n > best_len:
            best_base, best_len = a, n
    return best_base, best_len


def extract_numeric(scr):
    writes = _int_writes(scr)
    base, n = _longest_stride1_column([a for a, _ in writes])
    if base is None:
        return [], {}
    primary = set(range(base, base + n))
    records, buf = [], []
    for addr, varg in writes:
        buf.append((addr, varg))
        if addr in primary:               # primary write closes the record
            rid = addr - base
            fields = {f"0x{a - rid:x}": _val(v) for a, v in buf}
            records.append({"id": rid, "fields": fields})
            buf = []
    return records, {"primary_index_base": f"0x{base:x}", "record_span": n}


def extract_footer(scr):
    records = []
    for i, ins in enumerate(scr.instructions):
        if ins.opcode == COPY_LOCAL_ARRAY and ins.args and ins.args[0][0] == T_GLOBAL_INT:
            addr = ins.args[0][1]
            foff = ins.args[1][1]
            vals = read_footer_array(scr, foff)
            records.append({"id": i, "global_addr": f"0x{addr:x}",
                            "footer_off": f"0x{foff:x}",
                            "length": len(vals) if vals else 0,
                            "values": vals if vals else []})
    return records, {}


def join_messages(records: list[dict], message_scr) -> dict:
    """Join a message-dispatch script to INIT records by runtime id."""
    messages, message_meta = extract_message_table.extract_messages(message_scr)
    by_id = {message["id"]: message for message in messages}
    joined = 0
    for record in records:
        if message := by_id.get(record["id"]):
            record["message"] = {
                key: value for key, value in message.items() if key != "id"
            }
            joined += 1
    init_ids = {record["id"] for record in records}
    message_ids = set(by_id)
    return {
        "source": message_scr.path.name,
        **message_meta,
        "joined_count": joined,
        "init_ids_without_message": sorted(init_ids - message_ids),
        "message_ids_without_init": sorted(message_ids - init_ids),
    }


@cache
def _global_registry() -> dict:
    path = paths.BUILD / "globals.json"
    try:
        return json.loads(path.read_text(encoding="utf8")).get("globals", {})
    except (OSError, json.JSONDecodeError):
        return {}


def field_semantics(records: list[dict]) -> dict[str, str]:
    """Map raw extracted field keys to canonical semantic names when available."""
    keys = {
        key
        for record in records
        for key in (*record.get("fields", {}), *record.get("record_fields", {}))
    }
    registry = _global_registry()
    semantics = {}
    for key in sorted(keys, key=lambda value: tuple(
            int(part, 0) for part in value.split("/")
    )):
        parts = key.split("/")
        entry = registry.get(f"0x{int(parts[0], 16):x}", {})
        name = entry.get("name")
        if not name:
            continue
        if len(parts) == 3:
            column = parts[2]
            column_name = entry.get("columns", {}).get(column, f"column_{column}")
            name = f"{name}.{column_name}"
        semantics[key] = name
    return semantics


def write_data_index(data_dir: Path) -> None:
    """Regenerate the disposable build/data index from current table JSONs."""
    tables = []
    for path in sorted(data_dir.glob("*.json")):
        if path.name.endswith("-field-profile.json"):
            continue
        try:
            data = json.loads(path.read_text(encoding="utf8"))
        except (OSError, json.JSONDecodeError):
            continue
        if "table" in data and "record_count" in data:
            tables.append((path.name, data))
    lines = [
        "<!-- DO NOT EDIT -- generated by tools/extract_init.py -->",
        "# Parsed INIT data tables",
        "",
        "The JSON files in this directory are generated from SYS4 `*INIT` scripts. Raw global-array",
        "bases remain available in every record; confirmed field meanings live in",
        "`vm-map/globals.toml` and the generated `docs/global-reference.md`.",
        "",
        "| file | mode | records | messages | array fields | record columns |",
        "|---|---|---:|---:|---:|---:|",
    ]
    for filename, data in tables:
        columns = len(data.get("field_columns") or [])
        record_columns = len(data.get("record_field_columns") or [])
        message_count = data.get("message_table", {}).get(
            "joined_count", data.get("message_count", 0)
        )
        lines.append(
            f"| `{filename}` | {data['mode']} | {data['record_count']} | "
            f"{message_count} | {columns} | {record_columns} |"
        )
    lines += [
        "",
        "Name-mode tables expose one-based runtime `id` values, the lookup `name_array_base`,",
        "the first populated `name_write_base`, and the reserved `record_span`. Fields are keyed",
        "by the runtime lookup base used by `lookup-array`, not merely the first written cell.",
        "Linked row-major fields are stored separately in `record_fields`, keyed as",
        "`base/stride/column` from corpus-observed `lookup-array-2d` consumers.",
        "Where a matching `*MES` dispatcher exists, `message` preserves its player-facing",
        "title, description, furigana, and bytecode dispatch offset separately from the",
        "short description stored by the INIT script.",
        "Top-level `field_semantics` maps raw array/row-column keys to canonical machine-readable",
        "names from `vm-map/globals.toml`; raw keys remain intact as bytecode provenance.",
        "",
        "Use `tools/init_table_profile.py <TABLE> --build` to generate value/population and",
        "direct-consumer evidence. `STINIT` still requires a bespoke mixed numeric/string parser.",
        "",
    ]
    (data_dir / "README.md").write_text("\n".join(lines), encoding="utf8")


def main() -> int:
    argv = [a for a in sys.argv[1:] if not a.startswith("--")]
    mode_arg = next((sys.argv[i + 1] for i, a in enumerate(sys.argv) if a == "--mode"), None)
    if not argv:
        raise SystemExit(__doc__)
    name = argv[0].upper().removesuffix(".BIN")
    outname = argv[1] if len(argv) > 1 else name
    scr = sys4load.load(resolve(name))

    mode = mode_arg or detect_mode(scr)
    extractor = {"name": extract_name, "numeric": extract_numeric, "footer": extract_footer}[mode]
    recs, meta = extractor(scr)
    if mode == "name" and name in MESSAGE_TABLES:
        message_name = MESSAGE_TABLES[name]
        meta["message_table"] = join_messages(
            recs, sys4load.load(extract_message_table.resolve(message_name))
        )

    cols = sorted({c for r in recs for c in r.get("fields", {})}, key=lambda h: int(h, 16))
    semantics = field_semantics(recs)
    out = {"table": name, "source": scr.path.name, "magic": scr.magic, "mode": mode,
           "record_count": len(recs), **meta,
           "field_columns": cols if mode != "footer" else None,
           "field_semantics": semantics, "records": recs}
    outpath = paths.BUILD / "data" / f"{outname}.json"
    outpath.parent.mkdir(parents=True, exist_ok=True)
    outpath.write_text(json.dumps(out, ensure_ascii=False, indent=2), encoding="utf-8")
    write_data_index(outpath.parent)
    print(f"{name}: mode={mode}, {len(recs)} records"
          + (f", {len(cols)} field-columns" if mode != 'footer' else "")
          + f" -> build/data/{outname}.json")
    for r in recs[:4]:
        if mode == "footer":
            print(f"  id {r['id']:>4}  {r['global_addr']} <- footer {r['footer_off']} "
                  f"len {r['length']}  head={r['values'][:8]}")
        else:
            f4 = {k: r['fields'][k] for k in list(r['fields'])[:4]}
            print(f"  id {r['id']:>4}  {r.get('name','')!r:12} desc={r.get('desc','')!r} {f4}")
    return 0


if __name__ == "__main__":
    sys.exit(main())