import ctypes.util
import glob
import os
from pathlib import Path
from SCons.Script import Action, Value
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.system.camerad.cameras.nv12_info import get_nv12_info
from tools.laptop_device_build.qcom_artifact_contract import import_artifact


Import('env', 'arch')

# The CPU:LLVM compiler requires libLLVM; device and Metal builds do not.
if arch not in ('comma_arm64', 'Darwin') and ctypes.util.find_library('LLVM') is None:
  print("\033[93mmodeld: LLVM not found, skipping tinygrad model compilation. Install LLVM to compile models.\033[0m")
  Return()

lenv = env.Clone()
compiler_tree = (Dir('#tinygrad_model_runtime/legacy9d/tinygrad_repo')
                 if arch == 'comma_arm64' else Dir('#tinygrad_repo'))
lenv.PrependENVPath('PYTHONPATH', compiler_tree.abspath)
if arch == 'comma_arm64':
  lenv['ENV']['PYTHONDONTWRITEBYTECODE'] = '1'
  compiler_cache = Path(Dir('#.cache/model-compiler/9d0446a4ba8a532c8b674fb6ad795af015cd9dcf').abspath)
  lenv['ENV']['CACHEDB'] = str(compiler_cache / 'tinygrad.db')
  lenv['ENV']['XDG_CACHE_HOME'] = str(compiler_cache / 'xdg')

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

camera_configs = [(c.width, c.height) for c in (_ar_ox_fisheye, _os_fisheye)]
laptop_qcom_import = arch == 'comma_arm64' and os.getenv('COMMA_LAPTOP_CROSS_BUILD') == '1'
qcom_package = Dir('#.cache/qcom-artifact-package').abspath


def model_action(command):
  if not laptop_qcom_import:
    return Action(command, "  [ONNX] $TARGET")

  def do_import(target, source, env):
    try:
      import_artifact(Path(qcom_package), Path(Dir('#').abspath), Path(str(target[0])), command)
    except (OSError, ValueError, KeyError) as error:
      raise RuntimeError(f"Verified laptop QCOM package required in {qcom_package}: {error}") from error
  return Action(do_import, "  [QCOM IMPORT] $TARGET")


def warp_action(command):
  return model_action(command) if laptop_qcom_import else command

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.
  # Use Apple's clang for the CPU helper runtime without requiring libLLVM.
  tg_flags = 'DEV=METAL HCQ_RUNTIME_DEV=CPU:CLANG JIT=2' if arch == 'Darwin' else 'DEV=CPU:LLVM'

# 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 = str(Path(compiler_tree.abspath) / 'examples/openpilot')
# build on all cores
taskset = 'taskset -c 0-7 ' if arch == 'comma_arm64' else ''

def compile_model(onnx_path, pkl_path):
  onnx_path, target_pkl_path = File(onnx_path).abspath, File(pkl_path).abspath
  cmd = (f'{tg_flags} {mac_brew_string} {taskset}python3 "{compiler}/compile_onnx.py" '
         f'"{onnx_path}" "{target_pkl_path}" --out-of-band --benchmark-runs 1')
  node = lenv.Command(
    target_pkl_path,
    tinygrad_files + [onnx_path, Value(cmd)],
    model_action(cmd),
  )
  if laptop_qcom_import:
    lenv.AlwaysBuild(node)

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

model_w, model_h = MEDMODEL_INPUT_SIZE
for cam_w, cam_h in camera_configs:
  warp_pkl_path = File(f"models/driving_warp_{cam_w}x{cam_h}_tinygrad.pkl").abspath
  stride, y_height, uv_height, _ = get_nv12_info(cam_w, cam_h)
  cmd = (f'{tg_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}')
  node = lenv.Command(warp_pkl_path, tinygrad_files + warp_deps + [Value(cmd)], warp_action(cmd))
  if laptop_qcom_import:
    lenv.AlwaysBuild(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}')
  node = lenv.Command(dm_pkl_path, tinygrad_files + warp_deps + [Value(cmd)], warp_action(cmd))
  if laptop_qcom_import:
    lenv.AlwaysBuild(node)