import glob
import os
import shutil
import tempfile
import time
from SCons.Script import Action, Value
from openpilot.common.file_chunker import chunk_file, get_chunk_targets, get_existing_chunks, open_file_chunked
from openpilot.common.transformations.camera import _ar_ox_fisheye, _os_fisheye
from openpilot.common.transformations.model import MEDMODEL_INPUT_SIZE, DM_INPUT_SIZE
from openpilot.selfdrive.modeld.helpers import chestnut_present, modeld_pkl_path
from openpilot.system.camerad.cameras.nv12_info import get_nv12_info


Import('env', 'arch')
chunker_file = File("#openpilot/common/file_chunker.py")
lenv = env.Clone()
lenv.PrependENVPath('PYTHONPATH', Dir('#tinygrad_repo').abspath)

tinygrad_root = env.Dir("#").abspath
tinygrad_files = ["#"+x for x in glob.glob(env.Dir("#tinygrad_repo").relpath + "/**", recursive=True, root_dir=tinygrad_root)
                  if 'pycache' not in x and os.path.isfile(os.path.join(tinygrad_root, x))]

def estimate_pickle_max_size(onnx_size):
  # QCOM programs for models with spatial recurrent features can approach 2x
  # the ONNX size. Overestimating only adds an empty trailing chunk.
  return 2.0 * onnx_size + 10 * 1024 * 1024

camera_configs = [(c.width, c.height) for c in (_ar_ox_fisheye, _os_fisheye)]

if arch == 'comma_arm64':
  tg_flags = 'DEV=QCOM IMAGE=1 FLOAT16=1 NOLOCALS=1 JIT_BATCH_SIZE=0 OPENPILOT_HACKS=1'
else:
  # JIT=2 disables graph batching, which produces incorrect outputs after buffers change.
  tg_flags = 'DEV=METAL JIT=2' if arch == 'Darwin' else 'DEV=CPU:LLVM'

CHESTNUT = chestnut_present()
if CHESTNUT:
  chestnut_tg_flags = 'DEBUG=1 DEV=USB+AMD:LLVM FRAME_DEV=CPU FLOAT16=1 JIT_BATCH_SIZE=0 GMMU=0 TC_OPT=2 TC_MIN_GLOBALS=32'
  # the USB+AMD GPU takes an exclusive flock; serialize all targets that touch it
  chestnut_lock = File("models/.chestnut.lock").abspath

# tinygrad calls brew which needs a $HOME in the env
mac_brew_string = f'HOME={os.path.expanduser("~")}' if arch == 'Darwin' else ''

warp_deps = [File("#openpilot/system/camerad/cameras/nv12_info.py")]
compiler = Dir('#tinygrad_repo/examples/openpilot').abspath
# CPU 7 is isolated with isolcpus on AGNOS, so explicitly pin the compiler to it.
taskset = 'taskset -c 7 ' if arch == 'comma_arm64' else ''

def chestnut_action(command, pkl=None, chunks=()):
  def do_compile(target, source, env):
    from openpilot.system.hardware.chestnut.flash import link_up
    # chestnut can enumerate before its PCIe link is up due to varying 12V power behavior across cars
    for _ in range(10):
      if link_up():
        break
      time.sleep(1)
    else:
      print("Chestnut not ready, skipping big model build")
      return
    if ret := env.Execute(command):
      return ret
    if chunks:
      chunk_file(pkl, chunks)
  return Action(do_compile, "  [CHESTNUT] $TARGET")

def compile_model(onnx_path, pkl_path, flags, chestnut=False):
  onnx_path, target_pkl_path = File(onnx_path).abspath, File(pkl_path).abspath
  onnx_deps = get_existing_chunks(onnx_path)
  cmd = (f'{flags} {mac_brew_string} {taskset}python3 "{compiler}/compile_onnx.py" '
         f'"{{onnx}}" "{target_pkl_path}" --device-input "*" --out-of-band --benchmark-runs 1')
  def do_compile(target, source, env):
    if os.path.isfile(onnx_path):
      return env.Execute(cmd.format(onnx=onnx_path))
    # TODO: Remove ONNX chunk reassembly once models are precompiled.
    with tempfile.NamedTemporaryFile(dir=os.path.dirname(onnx_path), suffix='.onnx') as tmp, open_file_chunked(onnx_path) as src:
      shutil.copyfileobj(src, tmp)
      tmp.flush()
      return env.Execute(cmd.format(onnx=tmp.name))
  compile_action = Action(do_compile, "  [ONNX] $TARGET")
  onnx_sizes_sum = sum(os.path.getsize(f) for f in onnx_deps)
  chunk_targets = get_chunk_targets(target_pkl_path, estimate_pickle_max_size(onnx_sizes_sum))
  def do_chunk(target, source, env, pkl=target_pkl_path, chunks=chunk_targets):
    chunk_file(pkl, chunks)
  actions = chestnut_action(compile_action, target_pkl_path, chunk_targets) if chestnut else [compile_action, Action(do_chunk, "  [CHUNK] $TARGET")]
  node = lenv.Command(
    chunk_targets,
    tinygrad_files + onnx_deps + [Value(cmd), Value(chunk_targets), chunker_file],
    actions,
  )
  if chestnut:
    lenv.SideEffect(chestnut_lock, node)

compile_model('models/dmonitoring_model.onnx', 'models/dmonitoring_model_tinygrad.pkl', tg_flags)

model_w, model_h = MEDMODEL_INPUT_SIZE
for chestnut in [False, True] if CHESTNUT else [False]:
  file_prefix, cmd_flags = ('big_', chestnut_tg_flags) if chestnut else ('', tg_flags)
  compile_model(f'models/{file_prefix}driving_supercombo.onnx', modeld_pkl_path(chestnut), cmd_flags, chestnut)
  for cam_w, cam_h in camera_configs:
    warp_pkl_path = File(f"models/{file_prefix}driving_warp_{cam_w}x{cam_h}_tinygrad.pkl").abspath
    stride, y_height, uv_height, _ = get_nv12_info(cam_w, cam_h)
    cmd = (f'{cmd_flags} {mac_brew_string} {taskset}python3 "{compiler}/compile_warp.py" '
           f'--frame {cam_w},{cam_h},{stride},{y_height},{uv_height},{stride * (y_height + uv_height)} '
           f'--warp-to {model_w}x{model_h} --layout yuv420 --frames 2 '
           f'--output {warp_pkl_path}')
    action = chestnut_action(cmd) if chestnut else cmd
    node = lenv.Command(warp_pkl_path, tinygrad_files + warp_deps + [Value(cmd)], action)
    if chestnut:
      lenv.SideEffect(chestnut_lock, node)

dm_w, dm_h = DM_INPUT_SIZE
for cam_w, cam_h in camera_configs:
  dm_pkl_path = File(f"models/dm_warp_{cam_w}x{cam_h}_tinygrad.pkl").abspath
  stride, y_height, uv_height, frame_size = get_nv12_info(cam_w, cam_h)
  cmd = (f'{tg_flags} {mac_brew_string} python3 "{compiler}/compile_warp.py" '
         f'--frame {cam_w},{cam_h},{stride},{y_height},{uv_height},{frame_size} --warp-to {dm_w}x{dm_h} '
         f'--layout luma --border-fill 16 --transform-device NPY --output {dm_pkl_path}')
  lenv.Command(dm_pkl_path, tinygrad_files + warp_deps + [Value(cmd)], cmd)