#!/usr/bin/env python3
"""INPUT REPRESENTATION v2 — stop destroying information.

Phase 2 measured the tuned model reproducing docstrings 1/80 times, and I nearly filed it as
hallucination. It was my bug: the v1 disassembly emits ONLY the instruction stream, and a function's
docstring lives in `co_consts[0]` where NO opcode ever references it (CPython sets `__doc__`
implicitly at function-creation). The model was asked to reconstruct information it could not see.

WHAT IS ACTUALLY DESTROYED — measured (probe_consts.py), not assumed:
  * co_consts[0] docstrings: DESTROYED. Every one of them.
  * every other unreferenced const: only the implicit `None`. Carries nothing.
  * co_names: 0 unreferenced. Every name ALREADY appears in the instruction stream via argrepr
    (LOAD_GLOBAL/LOAD_ATTR/STORE_NAME...). Dumping a NAMES table would add zero information and
    cost input length -- so v2 does NOT dump it. (The brief asked for it; the measurement says no.)
  * signature shape: posonly/kwonly/vararg/kwarg boundaries were flattened into a flat arg list.
    Recoverable from co_posonlyargcount/co_kwonlyargcount/co_flags -- v1 just threw it away. This
    matters now that the corpus admits richer signatures than the pilot's.
  * default values: NOT destroyed (LOAD_CONST'd on the enclosing frame before MAKE_FUNCTION).

DESIGN — the rep must be a pure function of the code object, because at inference we have only a
.pyc. Everything below is read straight off the CodeType. Nothing is sourced from the .py file.

  CODE qualname(a, b, /, c, *args, d, **kw)     <- full signature shape
    DOC 'the docstring'                          <- co_consts[0], the destroyed information
    <instructions>
    EXC ...
  END

DOC is emitted as a first-class line rather than a raw CONSTS dump: a raw table would re-print
every already-referenced const (they appear inline at their LOAD_CONST) to recover exactly one
missing entry -- pure token inflation for no information. We emit precisely what was lost.
"""
from __future__ import annotations
import dis, types

JUMP_OPS = set(dis.hasjrel) | set(dis.hasjabs)

CO_VARARGS, CO_VARKEYWORDS = 0x04, 0x08


def _labels_for(co: types.CodeType) -> dict[int, str]:
    pts: set[int] = set()
    for i in dis.get_instructions(co):
        if i.opcode in JUMP_OPS and isinstance(i.argval, int):
            pts.add(i.argval)
    for e in dis._parse_exception_table(co):  # noqa: SLF001
        pts.update((e.start, e.end, e.target))   # e.end included — see the EXC note below
    return {off: f"L{n}" for n, off in enumerate(sorted(pts), 1)}


def _child_names(co: types.CodeType) -> dict[int, str]:
    return {id(c): c.co_qualname for c in co.co_consts if isinstance(c, types.CodeType)}


def _fmt_arg(i, labels: dict[int, str], kids: dict[int, str]) -> str:
    if i.arg is None:
        return ""
    if i.opcode in JUMP_OPS and isinstance(i.argval, int):
        return labels.get(i.argval, f"@{i.argval}")
    v = i.argval
    if isinstance(v, types.CodeType):
        return f"<code {kids.get(id(v), v.co_qualname)}>"
    if i.opname in ("LOAD_CONST", "RETURN_CONST", "KW_NAMES"):
        return repr(v)
    if i.argrepr:
        return i.argrepr
    return str(i.arg)


def signature(co: types.CodeType) -> str:
    """The full signature SHAPE, straight off the code object: posonly '/', kwonly '*',
    *args and **kwargs. v1 flattened all of this into a bare comma list."""
    n_pos, n_all = co.co_posonlyargcount, co.co_argcount
    n_kw = co.co_kwonlyargcount
    names = list(co.co_varnames)
    parts: list[str] = []
    idx = 0
    for k in range(n_all):
        parts.append(names[k])
        if k + 1 == n_pos and n_pos:
            parts.append("/")
    idx = n_all
    star_done = False
    if co.co_flags & CO_VARARGS:
        # *args sits right after the positional block
        va = names[n_all + n_kw]
        parts.append(f"*{va}")
        star_done = True
    if n_kw and not star_done:
        parts.append("*")
    for k in range(n_kw):
        parts.append(names[n_all + k])
    if co.co_flags & CO_VARKEYWORDS:
        off = n_all + n_kw + (1 if co.co_flags & CO_VARARGS else 0)
        parts.append(f"**{names[off]}")
    return ", ".join(parts)


def docstring_of(co: types.CodeType) -> str | None:
    """co_consts[0] iff it is the implicit docstring slot. This is the information v1 destroyed."""
    if co.co_consts and isinstance(co.co_consts[0], str):
        return co.co_consts[0]
    return None


def disassemble_v2(co: types.CodeType, out: list[str] | None = None) -> str:
    out = [] if out is None else out
    labels, kids = _labels_for(co), _child_names(co)
    out.append(f"CODE {co.co_qualname}({signature(co)})")
    doc = docstring_of(co)
    if doc is not None:
        out.append(f"  DOC {doc!r}")
    for i in dis.get_instructions(co):
        if i.offset in labels:
            out.append(f"{labels[i.offset]}:")
        a = _fmt_arg(i, labels, kids)
        out.append(f"  {i.opname} {a}".rstrip())
    for e in dis._parse_exception_table(co):  # noqa: SLF001
        # THE `end` IS LOAD-BEARING AND v2.0 OMITTED IT. Found at n=279: the model produced
        #     try: A            (except: pass) else: B; C
        # where the original was
        #     try: A; B; C      (except: pass)
        # CPython compiles an `else` block OUTSIDE the try range, so BOTH forms emit the SAME
        # instruction stream and differ ONLY in how far the exception table's range extends. With
        # `end` omitted, the two programs were BYTE-IDENTICAL in the model's input -- it could not
        # possibly tell them apart, and it guessed wrong. Exactly the docstring bug again: the model
        # was blamed for hallucinating information I had deleted.
        s = labels.get(e.start, f"@{e.start}")
        en = labels.get(e.end, f"@{e.end}")
        t = labels.get(e.target, f"@{e.target}")
        out.append(f"  EXC try={s}..{en} -> handler={t} depth={e.depth} lasti={e.lasti}")
    out.append("END")
    for c in co.co_consts:
        if isinstance(c, types.CodeType):
            disassemble_v2(c, out)
    return "\n".join(out)