import random, traceback, ctypes, argparse, os
from typing import Any
import numpy as np
from collections import defaultdict
from extra.optimization.helpers import load_worlds, ast_str_to_lin, kern_str_to_lin

# We need to insert ioctl before opening devices.
if os.getenv("VALIDATE_HCQ", 0) != 0:
  try:
    import extra.nv_gpu_driver.nv_ioctl
    from tinygrad import Device
    _, _ = Device["NV"], Device["CUDA"]
  except Exception: pass

  try:
    import extra.qcom_gpu_driver.opencl_ioctl
    from tinygrad import Device
    _, _ = Device["QCOM"], Device["GPU"]
  except Exception: pass

from tinygrad import Tensor, Device, dtypes
from tinygrad.tensor import _to_np_dtype
from tinygrad.codegen.kernel import Kernel
from tinygrad.codegen.kernel import Opt, OptOps
from tinygrad.engine.search import get_kernel_actions, bufs_from_lin
from tinygrad.engine.realize import CompiledRunner
from tinygrad.helpers import getenv, from_mv, prod, colored, Context, DEBUG, Timing
from tinygrad.ops import UOp, Ops
from tinygrad.device import is_dtype_supported

def on_linearizer_will_run(): pass
def on_linearizer_did_run(): pass
def compare_states(x, y): return (True, "")

if getenv("VALIDATE_HCQ"):
  if Device.DEFAULT == "NV":
    print("VALIDATE_HCQ: Comparing NV to CUDA")
    import extra.nv_gpu_driver.nv_ioctl
    validate_device = Device["CUDA"]
    on_linearizer_will_run = extra.nv_gpu_driver.nv_ioctl.before_launch
    on_linearizer_did_run = extra.nv_gpu_driver.nv_ioctl.collect_last_launch_state
    compare_states = extra.nv_gpu_driver.nv_ioctl.compare_launch_state
  elif Device.DEFAULT == "QCOM":
    print("VALIDATE_HCQ: Comparing QCOM to GPU")
    import extra.qcom_gpu_driver.opencl_ioctl
    validate_device = Device["GPU"]
    on_linearizer_will_run = extra.qcom_gpu_driver.opencl_ioctl.before_launch
    on_linearizer_did_run = extra.qcom_gpu_driver.opencl_ioctl.collect_last_launch_state
    compare_states = extra.qcom_gpu_driver.opencl_ioctl.compare_launch_state
  else:
    print(colored("VALIDATE_HCQ options is ignored", 'red'))

def tuplize_uops(uops:list[UOp]) -> tuple:
  return tuple([(x.op, x.dtype, tuple(uops.index(x) for x in x.src), x.arg) for x in uops])

device = Device[Device.DEFAULT]

def get_fuzz_rawbufs(lin):
  rawbufs = bufs_from_lin(lin)

  # Reallocate output buffer with additional area to detect out-of-bounds writes.
  RED_AREA_SIZE = 1024
  # setting output  # TODO: multi-output kernel
  rawbufs[0] = get_fuzz_rawbuf_like(rawbufs[0], zero=True, size=rawbufs[0].size+RED_AREA_SIZE)
  # setting inputs
  with Context(DEBUG=0):
    for rawbuf in rawbufs[1:]:
      if dtypes.is_unsigned(rawbuf.dtype):
        data = np.random.randint(0, 100, size=rawbuf.size, dtype=_to_np_dtype(rawbuf.dtype))
      elif dtypes.is_int(rawbuf.dtype):
        data = np.random.randint(-100, 100, size=rawbuf.size, dtype=_to_np_dtype(rawbuf.dtype))
      elif rawbuf.dtype == dtypes.bool:
        data = np.random.choice([True, False], size=rawbuf.size)
      elif rawbuf.dtype == dtypes.half:
        data = np.random.uniform(-1, 1, size=rawbuf.size).astype(dtype=_to_np_dtype(rawbuf.dtype))
      else:
        data = np.random.uniform(-10, 10, size=rawbuf.size).astype(dtype=_to_np_dtype(rawbuf.dtype))
      rawbuf.copyin(Tensor(data, device=lin.opts.device).realize().lazydata.base.realized.as_buffer())
  return rawbufs

def get_fuzz_rawbuf_like(old_rawbuf, zero=False, copy=False, size=None, force_device=None):
  rawbuf = type(old_rawbuf)(force_device or old_rawbuf.device, old_rawbuf.size if size is None else size, old_rawbuf.dtype).allocate()
  if copy:
    with Context(DEBUG=0): rawbuf.copyin(old_rawbuf.as_buffer())
  elif zero:
    with Context(DEBUG=0):
      mv = memoryview(bytearray(rawbuf.size * rawbuf.dtype.itemsize))
      ctypes.memset(from_mv(mv), 0, len(mv))
      rawbuf.copyin(mv)
  return rawbuf

def run_linearizer(lin: Kernel, rawbufs=None, var_vals=None) -> tuple[str, Any]: # (error msg, run state)
  if rawbufs is None: rawbufs = bufs_from_lin(lin)
  if var_vals is None: var_vals = {v: v.min for v in lin.vars}

  # TODO: images needs required_optimization
  try:
    prg = CompiledRunner(lin.to_program())
  except KeyboardInterrupt: raise
  except Exception:
    traceback.print_exc()
    return "COMPILE_ERROR", None

  if getenv("VALIDATE_HCQ"): on_linearizer_will_run()
  try:
    prg(rawbufs, var_vals, wait=True)
  except KeyboardInterrupt: raise
  except Exception:
    traceback.print_exc()
    return "EXEC_ERROR", None

  if getenv("VALIDATE_HCQ"): run_state = on_linearizer_did_run()
  else: run_state = None

  return "PASS", run_state

def compare_linearizer(lin: Kernel, rawbufs=None, var_vals=None, ground_truth=None, rtol=1e-2, atol=1e-2):
  # TODO: for bfloat16 it compiles linearizer, but it does not run because numpy cannot generate bf16 buffer.
  has_bf16 = any(b.dtype.base == dtypes.bfloat16 for b in lin.membufs)

  # TODO: raise specific fuzzing errors instead of str, and propagate the error message
  try:
    if rawbufs is None:
      rawbufs = get_fuzz_rawbufs(lin)
    else:
      rawbufs[0] = get_fuzz_rawbuf_like(rawbufs[0], zero=True) # get a new output buffer
  except KeyboardInterrupt: raise
  except BaseException:
    return ("RAWBUFS_ERROR", rawbufs, var_vals, ground_truth, None)

  if var_vals is None:
    # TODO: handle symbolic max case
    var_vals = {v: random.randint(v.vmin, v.vmax) for v in lin.ast.variables()}

  if ground_truth is None and not has_bf16:
    unoptimized = Kernel(lin.ast)
    unoptimized.required_optimizations()
    if run_linearizer(unoptimized, rawbufs, var_vals)[0] != "PASS":
      return ("BASELINE_ERROR", rawbufs, var_vals, ground_truth, None)
    ground_truth = np.frombuffer(rawbufs[0].as_buffer(), _to_np_dtype(rawbufs[0].dtype)).copy()

  rawbufs[0] = get_fuzz_rawbuf_like(rawbufs[0], zero=True) # get a new output buffer
  run_msg, run_state = run_linearizer(lin, rawbufs, var_vals)
  if run_msg != "PASS": return (run_msg, rawbufs, var_vals, ground_truth, run_state)

  try:
    if not has_bf16:
      result = np.frombuffer(rawbufs[0].as_buffer(), _to_np_dtype(rawbufs[0].dtype))
      np.testing.assert_allclose(result, ground_truth, rtol=rtol, atol=atol)
  except KeyboardInterrupt: raise
  except AssertionError as e:
    if DEBUG >= 2:
      print(f"COMPARE_ERROR details: {e}")
      if getenv("DEBUG_VALUES") > 0:
        mismatch_indices = np.where(~np.isclose(result, ground_truth, rtol=rtol, atol=atol))
        mismatched_result = result[mismatch_indices]
        mismatched_ground_truth = ground_truth[mismatch_indices]
        for i, idx in enumerate(mismatch_indices[0]):
          print(f"mismatch at {idx=}: result={mismatched_result[i]} <> ground_truth={mismatched_ground_truth[i]}")
    return ("COMPARE_ERROR", rawbufs, var_vals, ground_truth, run_state)

  return ("PASS", rawbufs, var_vals, ground_truth, run_state)

def fuzz_linearizer(lin: Kernel, rtol=1e-2, atol=1e-2, opts_list=None):
  SEED = getenv("SEED", 42)
  random.seed(SEED)
  np.random.seed(SEED)
  print(lin.ast)
  print(lin.colored_shape())
  seen_uops = {}
  last_lins = [lin]
  failures:defaultdict[str, list[tuple[tuple[UOp, ...], list[Opt]]]] = defaultdict(list)
  rawbufs, var_vals, ground_truth, validate_rawbufs = None, None, None, None

  FUZZ_ALL_ACTIONS = getenv("FUZZ_ALL_ACTIONS", 0)
  FUZZ_MAX_SIZE = getenv("FUZZ_MAX_SIZE", 0)
  FUZZ_IGNORE_SIMPLE_OPS = getenv("FUZZ_IGNORE_SIMPLE_OPS", 1)

  if FUZZ_MAX_SIZE > 0 and prod(lin.full_shape) > FUZZ_MAX_SIZE:
    print("skipping large kernel")
    return failures
  if FUZZ_IGNORE_SIMPLE_OPS and _is_simple(lin):
    print("skipping simple kernel")
    return failures

  test_depth = 1 if opts_list is not None else getenv("DEPTH", 1 if FUZZ_ALL_ACTIONS else 10)
  for depth in range(test_depth):
    next_lins = []
    for lin in last_lins:
      if opts_list is None: actions = get_kernel_actions(lin, include_0=False)
      else:
        actions = {}
        for oi,opts in enumerate(opts_list):
          lin2 = lin.copy()
          for o in opts: lin2.apply_opt(o)
          actions[oi] = lin2

      if not actions: continue
      if depth == 0 and getenv("FUZZ_REQUIRE_TC", 0):
        tc_acts = {i: k for k in actions.values() if k.applied_opts[0].op == OptOps.TC}
        if len(tc_acts) == 0: return failures
        else: actions = tc_acts

      test_lins = list(actions.values())
      if FUZZ_ALL_ACTIONS: print(f"testing {lin.applied_opts=} with {len(actions)} actions")
      elif opts_list is None: test_lins = [random.choice(test_lins)]

      for test_lin in test_lins:
        if not FUZZ_ALL_ACTIONS and test_lin.applied_opts: print(f"applied opts: {test_lin.applied_opts}")

        # stop if kernel uops repeat
        try: tuops = tuplize_uops(test_lin.linearize().uops)
        except KeyboardInterrupt: raise
        except BaseException as e:
          print(test_lin.ast)
          print(test_lin.applied_opts)
          print(e)
          failures["LINEARIZE_ERROR"].append((test_lin.ast, test_lin.applied_opts))
          continue

        if tuops in seen_uops: continue
        seen_uops[tuops] = tuple(test_lin.applied_opts)

        if not FUZZ_ALL_ACTIONS: print(test_lin.colored_shape())

        (msg, rawbufs, var_vals, ground_truth, state1) = compare_linearizer(test_lin, rawbufs, var_vals, ground_truth, rtol=rtol, atol=atol)
        if state1 is not None and validate_device is not None:
          validate_lin = test_lin.copy()
          validate_lin.opts = validate_device.renderer
          if validate_rawbufs is None:
            validate_rawbufs = [get_fuzz_rawbuf_like(x, copy=True, force_device=validate_device.device) for x in rawbufs]
          (_msg, _, _, _, state2) = compare_linearizer(validate_lin, validate_rawbufs, var_vals, ground_truth, rtol=rtol, atol=atol)

          if _msg != "PASS": failures[f"VALIDATE_DEV_{_msg}"].append((validate_lin.ast, validate_lin.applied_opts))

          ok, err_msg = compare_states(state1, state2)
          if not ok: failures["HCQ_COMPARE_FAILURE"].append((err_msg, test_lin.ast, test_lin.applied_opts, state1, state2))

        if msg != "PASS":
          print(test_lin.ast)
          print(test_lin.applied_opts)
          print(msg)
          failures[msg].append((test_lin.ast, test_lin.applied_opts))
          continue

        next_lins.append(test_lin)

    last_lins = next_lins
    if FUZZ_ALL_ACTIONS: print(f"depth={depth} total_lins={len(last_lins)} {failures=}")
  return failures

def _is_simple(lin: Kernel) -> bool:
  if len(lin.ast.src) > 1: return False
  ast:UOp = lin.ast.src[0]
  if ast.src[0].op is Ops.CAST and ast.src[0].src[0].op is Ops.LOAD: return True
  return False

if __name__ == "__main__":
  parser = argparse.ArgumentParser(description="Run a fuzz testing on one or more kernels", formatter_class=argparse.ArgumentDefaultsHelpFormatter)
  parser.add_argument("--ast", type=str, default=None, help="the ast for the kernel to be optimized")
  parser.add_argument("--file", type=str, default=None, help="a file containing asts to be optimized, one per line")
  parser.add_argument("--beamreplay", type=str, default=None, help="replay asts and opts got from beam with CAPTURE_BEAM")
  parser.add_argument("--logfile", type=str, default=None, help="a file containing a tuple of ast and applied_opts, one per line")
  parser.add_argument("--expected-failures", type=int, default=0, help="the number of expected failed kernels")
  parser.add_argument("--rtol", type=float, default=1e-2, help="relative tolerance for numerical comparison")
  parser.add_argument("--atol", type=float, default=1e-2, help="absolute tolerance for numerical comparison")
  args = parser.parse_args()

  opts_list = None
  if args.ast is not None:
    print("loaded AST from CLI")
    ast_strs = [args.ast]
  elif args.file is not None:
    print(f"loading ASTs from file '{args.file}'")
    with open(args.file, 'r') as file:
      ast_strs = file.readlines()
  elif args.beamreplay is not None:
    print(f"loading BEAM replay from file '{args.beamreplay}'")
    with open(args.beamreplay, 'r') as file: fdata = file.readlines()
    ast_strs, opts_list = [x.split(' :: ')[0] for x in fdata if not x.startswith("#")], [x.split(' :: ')[1] for x in fdata if not x.startswith("#")]

    # dedup ast_strs and opts_list
    dct = defaultdict(list)
    for i in range(len(ast_strs)): dct[ast_strs[i]].append(eval(opts_list[i]))
    ast_strs_items = list(dct.keys())
    opts_list = [dct[c] for c in ast_strs_items]
  elif args.logfile is not None:
    print(f"loading ASTs from LOGKERNS file '{args.file}'")
    with open(args.logfile, 'r') as file:
      kern_strs = file.readlines()
      test_lins = [kern_str_to_lin(kern_str) for kern_str in kern_strs]
      ast_strs = [f"{lin.ast}" for lin in test_lins]
  else:
    print("loading ASTs from world")
    ast_strs = load_worlds(filter_reduce=False, filter_novariable=False)

  print(f"{len(ast_strs)=}")
  tested = 0
  failed_ids = []
  failures = defaultdict(list)
  seen_ast_strs = set()

  try:
    for i, ast in enumerate(ast_strs[:getenv("FUZZ_N", len(ast_strs))]):
      if (nth := getenv("FUZZ_NTH", -1)) != -1 and i != nth: continue
      if getenv("FUZZ_IMAGEONLY") and "dtypes.image" not in ast: continue
      if "dtypes.image" in ast and Device.DEFAULT not in {"GPU", "QCOM"}: continue  # IMAGE is only for GPU
      if ast in seen_ast_strs: continue
      seen_ast_strs.add(ast)

      lin = ast_str_to_lin(ast)
      if not all(is_dtype_supported(buf.dtype) for buf in lin.bufs):
        print("skipping kernel due to not supported dtype")
        continue

      with Timing(f"tested ast {i}: "):
        tested += 1
        fuzz_failures = fuzz_linearizer(lin, rtol=args.rtol, atol=args.atol, opts_list=(opts_list[i] if opts_list else None))
        if fuzz_failures: failed_ids.append(i)
        for k, v in fuzz_failures.items():
          for f in v:
            failures[k].append(f)
  except KeyboardInterrupt: print(colored("STOPPING...", 'red'))

  for msg, errors in failures.items():
    for i, payload in enumerate(errors):
      print(f"{msg} {i} kernel: {payload}") # easier to use with output with verify_kernel.py

  print(f"{tested=}")
  if failures:
    print(f"{failed_ids=}")
    for msg, errors in failures.items():
      print(f"{msg}: {len(errors)}")
    if len(failed_ids) == args.expected_failures:
      print(colored(f"{len(failed_ids)} failed as expected", "yellow"))
  if len(failed_ids) != args.expected_failures:
    print(colored(f"failed on {len(failed_ids)} kernels, expected {args.expected_failures}", "red"))
    # TODO: fix this
    # raise RuntimeError(f"failed on {len(failed_ids)} kernels, expected {args.expected_failures}")
  else:
    print(colored("all passed", "green"))