"""
Copyright © IQ.Lvbs, apart of Project Teal Lvbs, All Rights Reserved, licensed under https://konn3kt.com/tos/
"""
from __future__ import annotations

import os
from dataclasses import dataclass
from pathlib import Path

import numpy as np
from tinygrad.nn.onnx import OnnxRunner
from tinygrad.tensor import Tensor

import iqpilot.selfdrive.iqmodeld.models.helpers as bundle_helpers
import iqpilot.selfdrive.iqmodeld.models.runners.model_runner as model_runner_mod
import iqpilot.selfdrive.iqmodeld.models.runners.tinygrad.tinygrad_runner as tinygrad_runner_mod
from iqpilot.selfdrive.iqmodeld.models.runners.model_runner import ModelType
from iqpilot.selfdrive.iqmodeld.models.runners.tinygrad.tinygrad_runner import TinygradRunner


SHARE_ROOT = Path(os.getenv("IQPILOT_SELECTOR_SHARE", "/Volumes/New New Vault/IQModels/models/recompiled16"))


@dataclass
class _TypeWrap:
  raw: int


@dataclass
class _Artifact:
  fileName: str


class _Model:
  def __init__(self, model_type: int, artifact_name: str, metadata_name: str):
    self.type = _TypeWrap(model_type)
    self.artifact = _Artifact(artifact_name)
    self.metadata = _Artifact(metadata_name)


class _Bundle:
  def __init__(self, models: list[_Model], is_20hz: bool = False):
    self.models = models
    self.is20hz = is_20hz


def _find_selector_dirs(limit: int = 3, require_onnx: bool = False) -> list[Path]:
  found: list[Path] = []
  if not SHARE_ROOT.is_dir():
    return found

  for bundle_dir in sorted(SHARE_ROOT.iterdir()):
    if not bundle_dir.is_dir():
      continue
    vision = next(bundle_dir.glob("driving_vision*_tinygrad.pkl"), None)
    policy = next(bundle_dir.glob("driving_policy*_tinygrad.pkl"), None)
    vision_meta = next(bundle_dir.glob("driving_vision*_metadata.pkl"), None)
    policy_meta = next(bundle_dir.glob("driving_policy*_metadata.pkl"), None)
    has_onnx = (bundle_dir / "driving_vision.onnx").is_file() and (bundle_dir / "driving_policy.onnx").is_file()
    if vision and policy and vision_meta and policy_meta and (has_onnx or not require_onnx):
      found.append(bundle_dir)
    if len(found) >= limit:
      break
  return found


def _seed_runner_inputs(runner: TinygradRunner) -> None:
  for name, shape in runner.input_shapes.items():
    runner.inputs[name] = Tensor(
      np.zeros(shape, dtype=np.float32),
      device=runner.input_to_device[name],
      dtype=runner.input_to_dtype[name],
    ).realize()


def _bundle_for_dir(bundle_dir: Path) -> _Bundle:
  vision = next(bundle_dir.glob("driving_vision*_tinygrad.pkl"))
  policy = next(bundle_dir.glob("driving_policy*_tinygrad.pkl"))
  vision_meta = next(bundle_dir.glob("driving_vision*_metadata.pkl"))
  policy_meta = next(bundle_dir.glob("driving_policy*_metadata.pkl"))
  return _Bundle([
    _Model(ModelType.vision, vision.name, vision_meta.name),
    _Model(ModelType.policy, policy.name, policy_meta.name),
  ])


def _run_tinygrad_bundle(bundle_dir: Path, monkeypatch):
  bundle = _bundle_for_dir(bundle_dir)
  monkeypatch.setattr(bundle_helpers, "get_active_bundle", lambda params=None: bundle, raising=False)
  monkeypatch.setattr(model_runner_mod, "get_active_bundle", lambda params=None: bundle, raising=False)
  monkeypatch.setattr(model_runner_mod, "_fetch_bundle", lambda: bundle)
  monkeypatch.setattr(tinygrad_runner_mod, "CUSTOM_MODEL_PATH", str(bundle_dir), raising=False)
  monkeypatch.setattr(model_runner_mod, "CUSTOM_MODEL_PATH", str(bundle_dir), raising=False)

  vision_runner = TinygradRunner(ModelType.vision)
  _seed_runner_inputs(vision_runner)
  vision_outputs = vision_runner.run_model()

  policy_runner = TinygradRunner(ModelType.policy)
  _seed_runner_inputs(policy_runner)
  policy_outputs = policy_runner.run_model()

  return vision_outputs, policy_outputs


def _run_onnx_bundle(bundle_dir: Path):
  vision_session = OnnxRunner(bundle_dir / "driving_vision.onnx")
  policy_session = OnnxRunner(bundle_dir / "driving_policy.onnx")

  def seed_inputs(session):
    seeded = {}
    for name, spec in session.graph_inputs.items():
      dtype_text = str(spec.dtype).lower()
      if "uchar" in dtype_text or "uint8" in dtype_text:
        seeded[name] = Tensor(np.zeros(spec.shape, dtype=np.uint8))
      elif "half" in dtype_text or "float16" in dtype_text:
        seeded[name] = Tensor(np.zeros(spec.shape, dtype=np.float16))
      else:
        seeded[name] = Tensor(np.zeros(spec.shape, dtype=np.float32))
    return seeded

  return (
    vision_session(seed_inputs(vision_session))["outputs"].numpy().flatten(),
    policy_session(seed_inputs(policy_session))["outputs"].numpy().flatten(),
  )


def test_three_selector_models_parse_via_share_onnx():
  if not SHARE_ROOT.is_dir():
    return
  selector_dirs = _find_selector_dirs(limit=3, require_onnx=True)
  assert len(selector_dirs) >= 3

  for bundle_dir in selector_dirs:
    vision_raw, policy_raw = _run_onnx_bundle(bundle_dir)
    assert vision_raw.size > 0
    assert policy_raw.size > 0


def test_selector_tinygrad_pkls_execute_when_host_compatible(monkeypatch):
  if not SHARE_ROOT.is_dir():
    return
  selector_dirs = _find_selector_dirs(limit=10)
  attempted = 0
  executed = 0

  for bundle_dir in selector_dirs:
    attempted += 1
    try:
      vision_outputs, policy_outputs = _run_tinygrad_bundle(bundle_dir, monkeypatch)
    except AssertionError as exc:
      if "Model was built on C3 or C3X" in str(exc):
        continue
      raise
    except FileNotFoundError as exc:
      if "/dev/kgsl-3d0" in str(exc):
        continue
      raise
    except TypeError as exc:
      if "DType.__init__()" in str(exc):
        continue
      raise

    assert "pose" in vision_outputs
    assert "plan" in policy_outputs
    executed += 1
    if executed >= 3:
      break

  if executed == 0:
    assert attempted > 0, "no selector bundles were inspected on the share"