import importlib.util
import base64
import json
import os
from pathlib import Path
import pickle
import subprocess
import sys
import tempfile
import unittest
from functools import partial
from unittest.mock import patch

import numpy as np

_script_path = Path(__file__).resolve().parents[4] / "scripts" / "model_compiler.py"
_script_spec = importlib.util.spec_from_file_location("starpilot_test_model_compiler", _script_path)
assert _script_spec is not None and _script_spec.loader is not None
model_compiler = importlib.util.module_from_spec(_script_spec)
_script_spec.loader.exec_module(model_compiler)


def varint(value):
  result = bytearray()
  while value > 127:
    result.append((value & 127) | 128)
    value >>= 7
  return bytes(result) + bytes([value])


def field(number, value):
  if isinstance(value, int):
    return varint(number << 3) + varint(value)
  if isinstance(value, str):
    value = value.encode()
  return varint((number << 3) | 2) + varint(len(value)) + value


def value_info(name, shape):
  dimensions = b"".join(field(1, field(1, dimension)) for dimension in shape)
  return field(1, name) + field(2, field(1, field(1, 1) + field(2, dimensions)))


def stateful_onnx():
  inputs = {"new_img": (2, 6, 2, 2), "state": (1, 2), "desire": (1, 2)}
  add = field(1, "state") + field(1, "desire") + field(2, "next_state") + field(4, "Add")
  identity = field(1, "next_state") + field(2, "outputs") + field(4, "Identity")
  graph = field(1, add) + field(1, identity) + field(2, "stateful-fixture")
  graph += b"".join(field(11, value_info(name, shape)) for name, shape in inputs.items())
  graph += field(12, value_info("outputs", (1, 2))) + field(12, value_info("next_state", (1, 2)))
  metadata = field(1, "output_slices") + field(2, base64.b64encode(pickle.dumps({"fixture": slice(0, 2)})))
  return field(1, 8) + field(7, graph) + field(8, field(2, 13)) + field(14, metadata)


class CompileTest(unittest.TestCase):
  def test_legacy_split_and_supercombo_history_survive_arena_round_trip(self):
    from tinygrad import Context, TinyJit
    from openpilot.starpilot.models.compile import (compile_jit, make_random_images, make_run_split_policy, make_run_supercombo,
                                                   make_split_input_queues, make_supercombo_input_queues, FAST_POLICY_INPUTS)
    image_shapes = {"input_imgs": (1, 12, 2, 2), "big_input_imgs": (1, 12, 2, 2)}
    policy_shapes = {"features_buffer": (1, 3, 2), "desire": (1, 2, 2), "traffic_convention": (1, 2)}
    def vision(inputs):
      values = [inputs[name].cast("float32").mean().reshape(1, 1) for name in image_shapes]
      return {"outputs": values[0].cat(values[1], dim=1)}
    def policy(inputs):
      return {"outputs": inputs["features_buffer"].sum(axis=1) + inputs["desire"].sum(axis=1) + inputs["traffic_convention"]}
    def supercombo(inputs):
      return {"outputs": policy(inputs)["outputs"] + vision(inputs)["outputs"]}
    with Context(DEV="CPU:LLVM", IMAGE=0, FLOAT16=0):
      for kind in ("split", "supercombo"):
        with self.subTest(kind=kind):
          if kind == "split":
            metadata = {"vision": {"input_shapes": image_shapes, "output_slices": {"hidden_state": slice(0, 2)}},
                        "policy": {"input_shapes": policy_shapes}}
            run = make_run_split_policy(vision, {"policy": policy}, metadata, ["policy"], 2)
            queues = partial(make_split_input_queues, image_shapes, policy_shapes, 2)
          else:
            metadata = {"model": {"input_shapes": image_shapes | policy_shapes}}
            run = make_run_supercombo(supercombo, metadata, 2)
            queues = partial(make_supercombo_input_queues, image_shapes | policy_shapes, 2)
          compile_jit(TinyJit(run, prune=True), partial(make_random_images, ["warped"], (2, 6, 2, 2)),
                      FAST_POLICY_INPUTS, queues, validation_runs=5)

  def test_mapped_stateful_image_queue_survives_jit_reload_and_reset(self):
    from types import SimpleNamespace
    from tinygrad import Context, Device, Tensor, TinyJit
    from tinygrad.dtype import dtypes
    from openpilot.starpilot.models.compile import compile_jit, make_run_stateful_supercombo, make_stateful_input_queues

    metadata = {"input_shapes": {"new_img": (2, 6, 4, 8), "action_t": (1, 2), "history": (1, 2)},
                "state_pairs": {"history": "next_history"}, "input_dtypes": {"history": "float32"},
                "mapped_image_input": True}

    class Model:
      graph_inputs = {"new_img": SimpleNamespace(dtype=dtypes.uint8), "action_t": SimpleNamespace(dtype=dtypes.float32)}

      def __call__(self, inputs):
        state = inputs["history"] + inputs["new_img"].float().mean() + inputs["action_t"]
        return {"outputs": state, "next_history": state}

    with Context(DEV="CPU:LLVM", IMAGE=0, FLOAT16=0):
      queues, _ = make_stateful_input_queues(metadata, Device.DEFAULT)
      image_buffer = queues["mapped_gpu_image"].uop.base.buffer
      self.assertTrue(image_buffer.options.cpu_access)
      self.assertIsNotNone(image_buffer.get_storage().host)
      run = make_run_stateful_supercombo(Model(), metadata)
      with self.assertRaisesRegex(ValueError, "Mapped stateful image queue is missing"):
        run(Tensor(np.zeros((2, 6, 4, 8), dtype=np.uint8)), queues["packed_npy_inputs"], history=queues["history"])
      compile_jit(TinyJit(run, prune=True),
                  lambda rng: {"warped": Tensor(rng.integers(0, 256, (2, 6, 4, 8), dtype=np.uint8), device="NPY").realize()},
                  ("packed_npy_inputs", "mapped_gpu_image", "history"),
                  partial(make_stateful_input_queues, metadata), validation_runs=3)

  def test_legacy_stateful_queue_layout_and_values_remain_unchanged(self):
    from tinygrad import Context, Device
    from openpilot.starpilot.models.compile import make_stateful_input_queues
    metadata = {"input_shapes": {"new_img": (2, 6, 4, 8), "action_t": (1, 2), "history": (1, 2)},
                "state_pairs": {"history": "next_history"}, "input_dtypes": {"history": "float32"}}
    with Context(DEV="CPU:LLVM"):
      queues, host = make_stateful_input_queues(metadata, Device.DEFAULT)
      self.assertEqual(tuple(queues), ("img_q", "big_img_q", "tfm", "big_tfm", "packed_npy_inputs", "history"))
      self.assertEqual(tuple(host), ("tfm", "big_tfm", "action_t"))
      expected_shapes = {"img_q": (1, 6, 4, 8), "big_img_q": (1, 6, 4, 8), "tfm": (3, 3), "big_tfm": (3, 3),
                         "packed_npy_inputs": (2,), "history": (1, 2)}
      for name, value in queues.items():
        with self.subTest(queue=name):
          self.assertEqual(value.shape, expected_shapes[name])
          self.assertFalse(value.uop.base.buffer.options.cpu_access)
          np.testing.assert_array_equal(value.numpy(), np.zeros(expected_shapes[name]))
      for value in host.values():
        np.testing.assert_array_equal(value, np.zeros(value.shape))

  def test_out_of_band_pickle_without_default_device_buffers(self):
    import io
    import struct
    from tinygrad import Context, Tensor
    from tinygrad_repo.examples.openpilot.helpers import dump_pickle, load_pickle

    with Context(DEV="CPU:LLVM"), tempfile.TemporaryDirectory() as temporary:
      path = Path(temporary) / "fixture.pkl"
      dump_pickle({"metadata": (1, "warp")}, path)
      self.assertEqual(load_pickle(path, out_of_band=True), {"metadata": (1, "warp")})
      values = np.arange(12, dtype=np.float32).reshape(2, 6)
      dump_pickle({"other_device": Tensor(values, device="NPY").realize()}, path)
      np.testing.assert_array_equal(load_pickle(path, out_of_band=True)["other_device"].numpy(), values)
      dump_pickle({"default_device": Tensor(values).realize()}, path)
      np.testing.assert_array_equal(load_pickle(path, out_of_band=True)["default_device"].numpy(), values)

      opcodes = io.BytesIO()
      marker = object()
      class ReferencingPickler(pickle.Pickler):
        def persistent_id(self, value):
          return (1, "invalid", 0) if value is marker else None
      pickler = ReferencingPickler(opcodes)
      pickler.dump(marker)
      path.write_bytes(struct.pack("<q", len(opcodes.getvalue())) + opcodes.getvalue())
      with self.assertRaises(pickle.UnpicklingError):
        load_pickle(path, out_of_band=True)

  def test_source_selection_rejects_ambiguity_and_missing_split_component(self):
    with tempfile.TemporaryDirectory() as temporary:
      root = Path(temporary)
      (root / "driving_supercombo.onnx").write_bytes(b"one")
      self.assertEqual(set(model_compiler.source_files(root, "test")), {"supercombo"})
      (root / "supercombo.onnx").write_bytes(b"two")
      with self.assertRaisesRegex(ValueError, "Ambiguous"):
        model_compiler.source_files(root, "test")
    with self.assertRaisesRegex(ValueError, "Split input requires"):
      model_compiler.driving_compile_args({"vision": Path("vision.onnx")}, "split")

  def test_multi_policy_order_and_explicit_behavior_version(self):
    kind, arguments = model_compiler.driving_compile_args({key: Path(key + ".onnx") for key in ("vision", "on-policy", "off-policy")}, "split")
    self.assertEqual(kind, "vision_multi_policy")
    self.assertEqual(arguments, ["--vision-onnx", "vision.onnx", "--on-policy-onnx", "on-policy.onnx", "--off-policy-onnx", "off-policy.onnx"])
    args = model_compiler.parse_args(["--example", "--version", "v15"])
    self.assertEqual((args.model, args.version, args.frame_skip), ("example", "v15", None))

  def test_failed_compile_preserves_existing_output(self):
    with tempfile.TemporaryDirectory() as temporary:
      root = Path(temporary)
      (root / "driving_supercombo.onnx").write_bytes(stateful_onnx())
      output = root / "test_driving_tinygrad.pkl"
      output.write_bytes(b"previous artifact")
      with patch.object(model_compiler.subprocess, "run", side_effect=subprocess.CalledProcessError(1, "compiler")):
        with self.assertRaises(subprocess.CalledProcessError):
          model_compiler.main(["--model", "test", "--version", "v16", "--input-dir", str(root), "--output-dir", str(root)])
      self.assertEqual(output.read_bytes(), b"previous artifact")

  def test_build_receipt_hashes_exact_sources_and_artifact(self):
    with tempfile.TemporaryDirectory() as temporary:
      root = Path(temporary)
      source = root / "driving_supercombo.onnx"
      source.write_bytes(stateful_onnx())
      child_run = subprocess.run
      scratch_paths = []
      def compile_stub(command, **kwargs):
        staged = Path(command[command.index("--output") + 1])
        result = child_run([sys.executable, "-c", "import tempfile; print(tempfile.mkdtemp(prefix='validation-'))"],
                           env=kwargs["env"], capture_output=True, text=True, check=True)
        scratch = Path(result.stdout.strip())
        self.assertEqual(scratch.parent, staged.parent)
        self.assertEqual(staged.parent.parent, root.resolve())
        scratch_paths.append(scratch)
        (scratch / "checkpoint").write_bytes(b"validation bytes")
        staged.write_bytes(b"compiled bytes")
      with patch.object(model_compiler.subprocess, "run", side_effect=compile_stub):
        model_compiler.main(["--model", "test", "--version", "v16", "--input-dir", str(root), "--output-dir", str(root)])
      output = root / "test_driving_tinygrad.pkl"
      receipt = json.loads(output.with_suffix(".build.json").read_text())
      self.assertEqual(receipt["source_sha256"], {source.name: model_compiler.sha256_file(source)})
      self.assertEqual(receipt["artifact_sha256"], model_compiler.sha256_file(output))
      self.assertEqual(receipt["artifact_abi"], "tinygrad_single_v1_arena")
      self.assertEqual(len(scratch_paths), 1)
      self.assertFalse(scratch_paths[0].parent.exists())

  def test_build_receipt_keeps_bank_compiler_when_global_runtime_advances(self):
    global_revision = "02e698160eed1751738212144a7ea2cfb160b3f1"
    bank_revision = "9d0446a4ba8a532c8b674fb6ad795af015cd9dcf"
    global_metadata = json.dumps({"dependencies": [{"path": "tinygrad_repo", "commit": global_revision}]})
    read_text = Path.read_text

    def read_with_advanced_global(path, *args, **kwargs):
      if path == model_compiler.REPO_ROOT / "upstream-sync.json":
        return global_metadata
      return read_text(path, *args, **kwargs)

    with tempfile.TemporaryDirectory() as temporary:
      root = Path(temporary)
      source = root / "driving_supercombo.onnx"
      source.write_bytes(stateful_onnx())
      compiled = b"compiled fixture bytes from pinned bank compiler"

      def compile_stub(command, **kwargs):
        self.assertEqual(command[:3], [sys.executable, "-m", "openpilot.starpilot.models.compile"])
        self.assertEqual(kwargs["cwd"], model_compiler.REPO_ROOT)
        staged = Path(command[command.index("--output") + 1])
        self.assertEqual(staged.parent.parent, root.resolve())
        staged.write_bytes(compiled)

      with patch.object(Path, "read_text", read_with_advanced_global), \
           patch.object(model_compiler.subprocess, "run", side_effect=compile_stub):
        # The old wrapper read this global metadata into its receipt. The bank
        # compiler remains pinned independently of the advanced global runtime.
        self.assertEqual(json.loads((model_compiler.REPO_ROOT / "upstream-sync.json").read_text())
                         ["dependencies"][0]["commit"], global_revision)
        model_compiler.main(["--model", "test", "--version", "v16", "--input-dir", str(root), "--output-dir", str(root)])
      output = root / "test_driving_tinygrad.pkl"
      receipt = json.loads(output.with_suffix(".build.json").read_text())
      self.assertEqual(output.read_bytes(), compiled)
      self.assertEqual(receipt["compiler_revision"], bank_revision)
      self.assertNotEqual(receipt["compiler_revision"], global_revision)
      self.assertEqual(receipt["source_sha256"], {source.name: model_compiler.sha256_file(source)})
      self.assertEqual(receipt["artifact_sha256"], model_compiler.sha256_file(output))
      self.assertEqual(receipt["artifact_size"], len(compiled))
      self.assertEqual(receipt["artifact_abi"], "tinygrad_single_v1_arena")

  def test_bank_compiler_rejects_same_origin_valid_revision_mismatch(self):
    from openpilot.starpilot.models import catalog
    from openpilot.starpilot.models.runtime_bridge import bootstrap, runtime

    global_revision = "02e698160eed1751738212144a7ea2cfb160b3f1"
    self.assertEqual(Path(catalog.__file__).resolve(),
                     (model_compiler.REPO_ROOT / "openpilot/starpilot/models/catalog.py").resolve())
    self.assertEqual(bootstrap.LEGACY_REVISION, runtime.LEGACY_REVISION)
    self.assertNotEqual(global_revision, bootstrap.LEGACY_REVISION)
    original_path = sys.path[:]
    original_environment = dict(os.environ)
    original_selection, original_protection = runtime._selection, runtime._protection

    def tinygrad_modules():
      return {name: module for name, module in sys.modules.items()
              if any(name == family or name.startswith(family + ".") for family in bootstrap.IMPORT_FAMILIES)}

    original_modules = tinygrad_modules()
    with patch.object(catalog, "COMPILER_REVISION", global_revision):
      with self.assertRaisesRegex(ValueError, "does not match the selected compiler runtime"):
        model_compiler.bank_compiler_revision()
    self.assertEqual(sys.path, original_path)
    self.assertEqual(dict(os.environ), original_environment)
    self.assertEqual(tinygrad_modules(), original_modules)
    self.assertIs(runtime._selection, original_selection)
    self.assertIs(runtime._protection, original_protection)

  def test_stateful_onnx_compiles_warp_and_replays_arena_on_cpu(self):
    with tempfile.TemporaryDirectory() as temporary:
      root = Path(temporary)
      source, artifact = root / "fixture.onnx", root / "fixture.pkl"
      source.write_bytes(stateful_onnx())
      environment = dict(os.environ, DEV="CPU:LLVM", JIT_BATCH_SIZE="0", IMAGE="0", FLOAT16="0")
      environment.pop("WARP_DEV", None)
      command = [sys.executable, "-m", "openpilot.starpilot.models.compile", "--model-type", "supercombo",
                 "--model-size", "4x4", "--camera-resolutions", "8x8", "--behavior-version", "v16",
                 "--supercombo-onnx", str(source), "--output", str(artifact)]
      result = subprocess.run(command, env=environment, capture_output=True, text=True, timeout=90)
      self.assertEqual(result.returncode, 0, result.stdout + result.stderr)
      validator = """import sys
import unittest
from pathlib import Path
from openpilot.starpilot.models.runtime_bridge.runtime import prepare_compiler
prepare_compiler()
import numpy as np
from tinygrad import Context
from tinygrad_repo.examples.openpilot.helpers import load_pickle
from openpilot.starpilot.models.compile import make_stateful_input_queues
artifact = Path(sys.argv[1])
self = unittest.TestCase()
with Context(DEV="CPU:LLVM"):
  loaded = load_pickle(artifact, out_of_band=True)
  queues, host = make_stateful_input_queues(loaded["metadata"]["model"], loaded["execution_device"])
  from tinygrad import Tensor
  warped = Tensor(np.zeros((2, 6, 2, 2), dtype=np.uint8)).realize()
  host["desire"][:] = (1, 2)
  for step in range(1, 4):
    output, = loaded["run_policy"](warped=warped, **{name: queues[name] for name in loaded["policy_input_keys"]})
    np.testing.assert_array_equal(output.numpy(), [[step, step * 2]])
  self.assertEqual(loaded["compiler_revision"], "9d0446a4ba8a532c8b674fb6ad795af015cd9dcf")
  self.assertEqual(loaded["frame_skip"], 1)
  self.assertIn((8, 8), loaded)
"""
      validation = subprocess.run([sys.executable, "-c", validator, str(artifact)], env=environment,
                                  capture_output=True, text=True, timeout=90)
      self.assertEqual(validation.returncode, 0, validation.stdout + validation.stderr)


if __name__ == "__main__":
  unittest.main()