"""
Generate Lib/_opcode_metadata.py for RustPython bytecode.
"""

import functools
import io
import itertools
import operator
import pathlib
import typing

from opcodes import OpcodeInfo
from utils import DEFAULT_INPUT, ROOT, get_conf, to_pascal_case, to_upper_snake_case

OUT_FILE = ROOT / "Lib/_opcode_metadata.py"


# Opcodes that needs to be first, regardless of their opcode ID.
PRIORITY_OPMAP = {
    "CACHE",
    "RESERVED",
    "RESUME",
    "INSTRUMENTED_LINE",
    "ENTER_EXECUTOR",
    "TRACE_RECORD",
}

INDENT = " " * 4


def main():
    override_conf = get_conf()
    inp = DEFAULT_INPUT.read_text()

    infos = tuple(OpcodeInfo.iter_infos(inp, override_conf))
    opcodes = tuple(itertools.chain.from_iterable(info.opcodes for info in infos))

    script_path = pathlib.Path(__file__).resolve().relative_to(ROOT).as_posix()
    out = io.StringIO()

    out.write(
        f"""
# This file is generated by {script_path}
# for RustPython bytecode format.
# Do not edit!
""".lstrip()
    )

    # _specializations
    out.write("_specializations = frozendict(\n")

    deopts = functools.reduce(operator.ior, map(operator.attrgetter("deopts"), infos))

    for key, lst in deopts.items():
        key = to_upper_snake_case(key)
        out.write(f"{INDENT}{key}=(\n")

        for item in map(to_upper_snake_case, dict.fromkeys(lst)):
            out.write(f'{INDENT * 3}"{item}",\n')

        out.write(f"{INDENT}),\n")

    out.write(")\n")

    # _specialized_opmap
    out.write("\n")
    out.write("_specialized_opmap = frozendict(\n")

    specialized = set(itertools.chain.from_iterable(deopts.values()))
    for opcode in sorted(opcodes, key=lambda op: op.cpython_name):
        if opcode.rust_name not in specialized:
            continue
        out.write(f"{INDENT}{opcode.cpython_name}={opcode.id},\n")

    out.write(")\n")

    # opmap
    out.write("\n")
    out.write("opmap = frozendict(\n")

    key = lambda op: (op.cpython_name not in PRIORITY_OPMAP, op.id)
    for opcode in sorted(opcodes, key=key):
        if opcode.rust_name in specialized:
            continue

        out.write(f"{INDENT}{opcode.cpython_name}={opcode.id},\n")

    out.write(")\n")

    # min
    out.write("\n")
    have_argument = min(opcode.id for opcode in opcodes if opcode.have_argument) - 1
    out.write(f"HAVE_ARGUMENT = {have_argument}\n")

    min_instrumented = min(opcode.id for opcode in opcodes if opcode.is_instrumented)
    out.write(f"MIN_INSTRUMENTED_OPCODE = {min_instrumented}\n")

    # write output
    generated = out.getvalue()
    OUT_FILE.write_text(generated)


if __name__ == "__main__":
    main()