import functools
import inspect
from typing import Any, Callable, Optional

from .translator import CudaEmitter
from .builder import CudaSourceBuilder


class KernelHandle:
  def __init__(self, func: Callable, builder: CudaSourceBuilder, name: str):
    self.func = func
    self.builder = builder
    self.name = name
    self._source: Optional[str] = None

  def get_source(self) -> str:
    if self._source is None:
      self._source = self.builder.build_source()
    return self._source

  def get_kernel_source(self) -> str:
    return self.builder.kernels.get(self.name, "")

  def write_to_disk(self, output_dir: str, base_name: Optional[str] = None, write_header: bool = True):
    bn = base_name or self.name
    return self.builder.write_to_disk(output_dir, bn, write_header=write_header)

  def __repr__(self):
    return f"<KernelHandle name={self.name!r}>"


class LaunchSpec:
  def __init__(self, handle: KernelHandle, grid, meta: dict):
    self.handle = handle
    self.grid = grid
    self.meta = meta

  def __call__(self, *args):
    try:
      from .cupy_bindings import CupyKernel, HAS_CUPY
      if HAS_CUPY:
        block = self.meta.get("block", (1, 1, 1))
        if isinstance(block, int):
          block = (block, 1, 1)
        ck = CupyKernel.from_builder(
          self.handle.builder,
          self.handle.name,
          self.grid,
          block,
        )
        ck(*args)
    except Exception as e:
      raise RuntimeError(f"Kernel launch failed: {e}") from e


class tritoncu_jit:
  def __init__(
    self,
    fn: Optional[Callable] = None,
    *,
    do_not_specialize: Optional[list] = None,
    debug: bool = False,
    noinline: bool = False,
  ):
    self._fn = fn
    self._do_not_specialize = do_not_specialize or []
    self._debug = debug
    self._noinline = noinline
    self._cache: dict = {}

    if fn is not None:
      self.__doc__ = fn.__doc__
      self.__wrapped__ = fn

  def __call__(self, fn: Optional[Callable] = None, **kwargs):
    if callable(fn) and self._fn is None:
      self._fn = fn
      self.__doc__ = fn.__doc__
      self.__wrapped__ = fn
      return self
    if self._fn is not None:
      return self._make_launch_spec(fn, kwargs) if fn else self._direct_call(**kwargs)
    raise RuntimeError("No function registered")

  def _direct_call(self, **kwargs):
    raise RuntimeError("Use [grid](...) syntax to launch kernels")

  def __getitem__(self, grid):
    return _GridBound(self, grid)

  def compile(self, meta: Optional[dict] = None) -> KernelHandle:
    meta = meta or {}
    cache_key = tuple(sorted(meta.items()))
    if cache_key in self._cache:
      return self._cache[cache_key]

    constexprs = {}
    emitter = CudaEmitter(self._fn, constexprs, meta)
    kernel_src = emitter.translate()
    builder = CudaSourceBuilder()
    builder.add_kernel(self._fn.__name__, kernel_src)
    handle = KernelHandle(self._fn, builder, self._fn.__name__)

    self._cache[cache_key] = handle
    return handle

  @property
  def __name__(self):
    return self._fn.__name__ if self._fn else "unknown"

  def __repr__(self):
    return f"<tritoncu_jit {self.__name__!r}>"


class _GridBound:
  def __init__(self, kernel: tritoncu_jit, grid):
    self.kernel = kernel
    self.grid = grid

  def __call__(self, *args, **meta):
    handle = self.kernel.compile(meta)
    spec = LaunchSpec(handle, self.grid, meta)
    spec(*args)
    return spec


def jit(
  fn: Optional[Callable] = None,
  *,
  do_not_specialize: Optional[list] = None,
  debug: bool = False,
  noinline: bool = False,
) -> tritoncu_jit:
  if fn is not None:
    return tritoncu_jit(fn, do_not_specialize=do_not_specialize, debug=debug, noinline=noinline)
  def decorator(f: Callable) -> tritoncu_jit:
    return tritoncu_jit(f, do_not_specialize=do_not_specialize, debug=debug, noinline=noinline)
  return decorator
