import inspect
import os
import sys
import platform
from functools import wraps, singledispatchmethod, partial
from itertools import product
from collections.abc import Callable, Sequence, Iterable
import packaging.version
import contextlib

import pytest
import torch
from torch._dynamo import is_inductor_supported
from torch.testing import assert_close

from lightning_utilities.core.imports import package_available

from thunder.core.pytree import tree_flatten, tree_map
import thunder.core.dtypes as datatypes
import thunder.core.devices as devices
import thunder.executors as executors
import thunder.extend as extend
import thunder.executors.triton_utils as triton_utils
import thunder.core.utils as utils

from thunder.core.trace import TraceCtx, detached_trace
from thunder.dynamo import thunderfx

import thunder

__all__ = [
    "TestExecutor",
    "nvFuserExecutor",
    "TorchExecutor",
]


# A marker for actually wanting NOTHING instead of an unspecified value (marked with None)
class NOTHING:
    pass


# Require Triton version 2.1 or greater, since our current Triton executor won't run
#   properly due to an error in 2.0
TRITON_AVAILABLE: bool = triton_utils.is_triton_version_at_least("2.1")

NVFUSER_AVAILABLE = executors.nvfuser_available()

IN_CI: bool = os.getenv("CI", None) == "true"
CUDA_AVAILABLE: bool = torch.cuda.is_available()
env_var_FORCE_CPU_TEST_INSTANTIATION: str = os.getenv("FORCE_CPU_TEST_INSTANTIATION", None)
FORCE_CPU_TEST_INSTANTIATION: bool = (
    env_var_FORCE_CPU_TEST_INSTANTIATION == "true" or env_var_FORCE_CPU_TEST_INSTANTIATION == "1"
)
env_var_DISABLE_CUDA_TEST_INSTANTIATION: str = os.getenv("DISABLE_CUDA_TEST_INSTANTIATION", None)
DISABLE_CUDA_TEST_INSTANTIATION: bool = (
    env_var_DISABLE_CUDA_TEST_INSTANTIATION == "true" or env_var_DISABLE_CUDA_TEST_INSTANTIATION == "1"
)
IS_WINDOWS = platform.system() == "Windows"


def _bitsandbytes_available():
    if not package_available("bitsandbytes"):
        return False
    try:
        import bitsandbytes

        if torch.cuda.is_available():
            import bitsandbytes.diagnostics.main

            bitsandbytes.diagnostics.main.sanity_check()
    except (ImportError, RuntimeError):
        return False
    return True


BITSANDBYTES_AVAILABLE = _bitsandbytes_available()


# TODO This change should be handled properly, this is a temporary fix to allow the CI to progress.
# See https://github.com/Lightning-AI/lightning-thunder/issues/2807
def _pytorch_removed_args_tensor_mask() -> bool:
    """Check if PyTorch removed args_tensor_mask from autograd_function_apply.

    PyTorch removed args_tensor_mask in https://github.com/pytorch/pytorch/pull/166788
    (commit 5cf15aef144fd03379a5796e37698eee5e4575b8, Dec 12, 2025).

    Returns True if args_tensor_mask is NO LONGER accepted (new PyTorch).
    """
    import re

    version = torch.__version__
    # Nightly versions look like: 2.6.0.dev20251212+cpu
    match = re.search(r"\.dev(\d{8})", version)
    if match:
        nightly_date = int(match.group(1))
        # args_tensor_mask removed around Dec 12, 2025
        return nightly_date >= 20251212

    # Lightning AI nightly builds have version strings like: 2.10.0a0+git62c80e7
    # We conservatively check against the base version (the portion before "+")
    # because we do not know whether the problematic commit removing args_tensor_mask
    # is before or after 2.10.0a0+git62c80e7. Therefore, we return True (masked as removed)
    # if we are on version 2.10.0a0 or later.
    base_version = packaging.version.parse(version.split("+")[0])
    if base_version >= packaging.version.parse("2.10.0a0"):
        return True

    return False


xfail_if_args_tensor_mask_removed = pytest.mark.xfail(
    _pytorch_removed_args_tensor_mask(),
    reason="PyTorch >= 2.10.0a0+git62c80e7 or nightly >= 20251212 removed args_tensor_mask from autograd_function_apply (PR #166788)",
)


def version_between(version: str, *, min_ver: str | None = None, max_ver: str | None = None):
    v = packaging.version.parse(version)
    if min_ver is not None and v < packaging.version.parse(min_ver):
        return False
    if max_ver is not None and v > packaging.version.parse(max_ver):
        return False
    return True


# Filters the CPU devicetype when in CI, CUDA is available, and the environment variable
#   FORCE_CPU_TEST_INSTANTIATION isn't forcing CPU test instantiation
def filter_ci_devicetypes(devicetypes: Iterable[devices.DeviceType]) -> tuple[devices.DeviceType]:
    filtered: tuple[devices.DeviceType]
    if IN_CI and CUDA_AVAILABLE and not FORCE_CPU_TEST_INSTANTIATION:
        filtered = tuple(x for x in devicetypes if x is not devices.DeviceType.CPU)
    else:
        filtered = tuple(x for x in devicetypes)

    if DISABLE_CUDA_TEST_INSTANTIATION:
        filtered = tuple(x for x in devicetypes if x is not devices.DeviceType.CUDA)

    return filtered


# Asserts that a candidate is closer to a reference than a competitor
# This is useful when trying to compare low precision operators; the
#   reference is typically the result in a higher precision datatype, like double,
#   while the candidate and competitor are often in bfloat16 or float16
def assert_closer(*, reference, candidate, competitor, comparator):
    def _to_meta(x):
        if isinstance(x, torch.Tensor):
            return x.to(device="meta")

        return x

    # Validates metadata
    reference_meta = tree_map(_to_meta, reference)
    candidate_meta = tree_map(_to_meta, candidate)
    competitor_meta = tree_map(_to_meta, competitor)

    assert_close(reference_meta, candidate_meta, check_dtype=False)
    assert_close(reference_meta, competitor_meta, check_dtype=False)
    comparator(candidate_meta, competitor_meta)

    reference_flats, _ = tree_flatten(reference)
    candidate_flats, _ = tree_flatten(candidate)
    competitor_flats, _ = tree_flatten(competitor)

    for ref, cand, com in zip(reference_flats, candidate_flats, competitor_flats):
        if isinstance(ref, torch.Tensor):
            candidate_dist = torch.abs(ref - cand)
            competitor_dist = torch.abs(ref - com)
            minimum_dist = torch.minimum(candidate_dist, competitor_dist)

            signed_minimum_dist = torch.where(ref > cand, -minimum_dist, minimum_dist)
            target = ref + signed_minimum_dist

            comparator(cand, target, check_dtype=False)


# TODO: Add device type functionality to an object in this list
def _all_devicetypes() -> Sequence[devices.DeviceType]:
    return devices.all_devicetypes


# TODO Technically CUDA can be available without a CUDA device and that might be interesting to test
def available_devicetypes():
    if torch.cuda.is_available() and torch.cuda.device_count() > 0:
        return devices.all_devicetypes
    return (devices.DeviceType.CPU,)


class TestExecutor:
    def is_available(self) -> bool:
        return True

    def supports_dtype(self, dtype: datatypes.dtype) -> bool:
        return dtype in datatypes.resolve_dtypes(self.supported_dtypes)

    def supports_devicetype(self, devicetype: devices.DeviceType) -> bool:
        return devicetype in self.supported_devicetypes

    # NOTE This method should be overridden by subclasses
    def executors_list(self) -> list[extend.Executor]:
        return []

    @singledispatchmethod
    def make_callable(self, fn, **kwargs):
        return thunder.jit(fn, executors=self.executors_list(), **kwargs)

    @make_callable.register
    def make_callable_from_trace(self, trace: TraceCtx, **kwargs):
        # transform_for_execution doesn't work without a set trace
        # So we use detached_trace to get the tracectx and then use it
        with detached_trace():
            traces = thunder.common.transform_for_execution(trace, executors_list=self.executors_list(), **kwargs)
        return traces[-1].python_callable()


# TODO Convert to singletons or just add to executor logic
class nvFuserTestExecutor(TestExecutor):
    name = "nvfuser"
    supported_devicetypes = (devices.DeviceType.CUDA,)
    supported_dtypes = (
        *datatypes.float_math_dtypes,
        datatypes.bool8,
        datatypes.int32,
        datatypes.int64,
        datatypes.complex64,
        datatypes.complex128,
    )

    def executors_list(self) -> list[extend.Executor]:
        return [executors.get_nvfuser_executor()]

    def version(self):
        return executors.get_nvfuser_executor().version


# TODO Convert to singletons or just add to executor logic
class TorchTestExecutor(TestExecutor):
    name = "torch"
    supported_devicetypes = (devices.DeviceType.CPU, devices.DeviceType.CUDA)
    supported_dtypes = (datatypes.dtype,)

    def executors_list(self) -> list[extend.Executor]:
        return [executors.get_torch_executor()]

    def version(self):
        return torch.__version__


class TorchCompileXentropyTestExecutor(TestExecutor):
    name = "torchcompile_xentropy"
    supported_devicetypes = (devices.DeviceType.CPU, devices.DeviceType.CUDA)
    supported_dtypes = (datatypes.dtype,)

    def is_available(self) -> bool:
        return not IS_WINDOWS

    def executors_list(self) -> list[extend.Executor]:
        from thunder.executors.torch_compile import torch_compile_cat_ex

        return [torch_compile_cat_ex]

    def version(self):
        return torch.__version__


class TorchCompileCatTestExecutor(TestExecutor):
    name = "torchcompile_cat"
    supported_devicetypes = (devices.DeviceType.CPU, devices.DeviceType.CUDA)
    supported_dtypes = (datatypes.dtype,)

    def is_available(self) -> bool:
        return not IS_WINDOWS

    def executors_list(self) -> list[extend.Executor]:
        from thunder.executors.torch_compile import torch_compile_cat_ex

        return [torch_compile_cat_ex]

    def version(self):
        return torch.__version__


class TorchCompileTestExecutor(TestExecutor):
    name = "torchcompile"
    supported_devicetypes = (devices.DeviceType.CPU, devices.DeviceType.CUDA)
    supported_dtypes = (datatypes.dtype,)

    def is_available(self) -> bool:
        return not IS_WINDOWS

    def executors_list(self) -> list[extend.Executor]:
        from thunder.executors.torch_compile import torch_compile_ex

        return [torch_compile_ex]

    def version(self):
        return torch.__version__


# This is a test executor that uses the ThunderCompiler with torch.compile to
# compile the function. However, this executor is not used for all tests by
# default (it's not part of the _all_test_executors list) because it might
# increase the test runtime significantly. Instead, it's used for specific tests
# that add it to the supported_executors list when needed. Thunder's end-to-end
# tests (test_networks.py) use this executor to test the integration between
# torch.compile and Thunder.
class DynamoThunderTestExecutor(TestExecutor):
    name = "DynamoThunder"
    supported_devicetypes = (devices.DeviceType.CPU, devices.DeviceType.CUDA)
    supported_dtypes = (datatypes.dtype,)

    def make_callable(self, fn, **kwargs):
        return thunderfx(fn, **kwargs)


# TODO Refactor these executors into the actual executor (sub)modules
TorchExecutor: TorchTestExecutor = TorchTestExecutor()
TorchCompileCatExecutor: TorchCompileCatTestExecutor = TorchCompileCatTestExecutor()
TorchCompileXentropyExecutor: TorchCompileXentropyTestExecutor = TorchCompileXentropyTestExecutor()
TorchCompileExecutor: TorchCompileTestExecutor = TorchCompileTestExecutor()
DynamoThunderExecutor: DynamoThunderTestExecutor = DynamoThunderTestExecutor()
nvFuserExecutor: None | nvFuserTestExecutor = None

if NVFUSER_AVAILABLE:
    nvFuserExecutor = nvFuserTestExecutor()


def _all_test_executors():
    """Constructs a list of all Thunder executors to be used when generating tests."""
    # TODO: include the torch compile executors: https://github.com/Lightning-AI/lightning-thunder/issues/299
    executors = [TorchExecutor]

    if NVFUSER_AVAILABLE:
        executors.append(nvFuserExecutor)

    return executors


# Translates test templates with names like test_foo into instantiated tests with names like
#   test_foo_nvFuser_CUDA_float32
# TODO Fix test name when dtype is None
# TODO Refactor with _instantiate_opinfo_test_template
# TODO Should the device str include the actual device number, or just the device type like now?
# TODO Support multiple devices
def _instantiate_executor_test_template(
    template: Callable,
    scope,
    *,
    executor: TestExecutor,
    device_or_devices: devices.Device | Sequence[devices.Device],
    dtype: datatypes.dtype,
    as_name: str | None = None,
) -> Callable:
    devicetype: devices.DeviceType
    device_str: str | list[str]
    devicetype = device_or_devices
    if isinstance(device_or_devices, devices.Device):
        devicetype = device_or_devices.devicetype
        device_str = device_or_devices.device_str()
    else:
        devicetype = device_or_devices[0].devicetype
        device_str = []
        for device in device_or_devices:
            device_str.append(device.device_str())

    devicetype_str = devices.devicetype_string(devicetype)
    template_name = as_name if as_name is not None else template.__name__
    test_name = "_".join((template_name, executor.name, devicetype_str, str(dtype)))

    test = partial(template, executor, device_str, dtype)

    # Mimics the instantiated test
    # TODO Review this mimicry -- are there other attributes to mimic?
    test.__name__ = test_name
    test.__module__ = test.__module__

    return test


# TODO Support multiple devices
def _instantiate_opinfo_test_template(
    template: Callable, scope, *, opinfo, executor: TestExecutor, device: devices.Device, dtype: datatypes.dtype
) -> Callable:
    """Instantiates a test template for an operator."""

    device_str = devices.devicetype_string(device.devicetype)

    test_name = "_".join((template.__name__, opinfo.name, executor.name, device_str, str(dtype)))

    # Acquires the comparator
    # TODO If multiple decorators define custom comparators and they "overlap", then which
    #   custom comparator is applied by this is uncertain
    comp = assert_close
    for decorator in opinfo.test_decorators(template.__name__, executor, device, dtype):
        if isinstance(decorator, custom_comparator):
            comp = decorator.comparator

    def test():
        result = template(opinfo, device_str, dtype, executor, comp)
        return result

    # Applies decorators
    for decorator in opinfo.test_decorators(template.__name__, executor, device, dtype):
        test = decorator(test)

    # Mimics the instantiated test
    # TODO Review this mimicry -- are there other attributes to mimic?
    #   Probably want to refactor this to codeutils
    test.__name__ = test_name
    test.__module__ = test.__module__

    return test


# TODO Add documentation and example uses; not this must be the LAST decorator applied
# TODO Support more dtype specification flexibility
# TODO Add ability to run on other devices by default (like cuda:1 instead of cuda:0 being the default)
class ops:
    def __init__(
        self, opinfos, *, supported_executors=None, supported_devicetypes=None, supported_dtypes=None, scope=None
    ):
        self.opinfos = opinfos

        self.supported_executors = (
            set(supported_executors)
            if supported_executors is not None
            else set(_all_test_executors() + [TorchCompileCatExecutor, TorchCompileXentropyExecutor])
        )
        for ex in self.supported_executors:
            assert isinstance(ex, TestExecutor)

        self.supported_devicetypes = (
            set(supported_devicetypes) if supported_devicetypes is not None else set(_all_devicetypes())
        )
        self.supported_devicetypes = set(filter_ci_devicetypes(self.supported_devicetypes))

        # TODO: add support for float4_e2m1fn_x2 and float_8bit_dtypes
        skip_dtypes = datatypes.float_8bit_dtypes | {datatypes.float4_e2m1fn_x2, datatypes.float4_e2m1fn_x2_}

        self.supported_dtypes = (
            datatypes.resolve_dtypes(supported_dtypes)
            if supported_dtypes is not None
            else datatypes.all_dtypes - skip_dtypes
        )

        if supported_dtypes == NOTHING:
            self.supported_dtypes = NOTHING

        # Acquires the caller's global scope
        if scope is None:
            previous_frame = inspect.currentframe().f_back
            scope = previous_frame.f_globals
        self.scope = scope

    def __call__(self, test_template):
        # NOTE Unlike a typical decorator, this __call__ does not return a function, because it may
        #   (and typically does) instantiate multiple functions from the template it consumes
        #   Since Python doesn't natively support one-to-many function decorators, the produced
        #   functions are directly assigned to the requested scope (the caller's global scope by default)

        for opinfo in self.opinfos:
            devicetypes = (
                opinfo.devicetypes().intersection(self.supported_devicetypes).intersection(set(available_devicetypes()))
            )

            for executor, devicetype in product(
                sorted(self.supported_executors, key=lambda x: repr(x)), sorted(devicetypes, key=lambda x: repr(x))
            ):
                if not executor.supports_devicetype(devicetype):
                    continue

                if any("torchcompile" in ex.name for ex in executor.executors_list()) and (
                    not opinfo.test_torch_compile_executor or not is_inductor_supported()
                ):
                    continue

                device = devices.Device(devicetype, None)

                # TODO Pass device_type to dtypes()
                dtypes = opinfo.dtypes()
                if self.supported_dtypes != (None,):
                    dtypes = dtypes.intersection(self.supported_dtypes)

                for dtype in sorted(dtypes, key=lambda t: repr(t)):
                    if not executor.supports_dtype(dtype):
                        continue

                    test = _instantiate_opinfo_test_template(
                        test_template,
                        self.scope,
                        opinfo=opinfo,
                        executor=executor,
                        device=device,
                        dtype=dtype,
                    )
                    # Adds the instantiated test to the requested scope
                    self.scope[test.__name__] = test

                    # [NOTE] dynamo reset
                    # dynamo caches are per code object, not frame. All the dynamic copies of these context
                    # managers share the same code cache in the process. This is a problem in a single pytest process
                    # that runs many traces
                    if any("torchcompile" in ex.name for ex in executor.executors_list()):
                        torch._dynamo.reset()


# TODO Allow executing the test suite on different devices (not just always cuda:0)
# TODO Example uses, note this must be the LAST decorator applied
class instantiate:
    # TODO: support other kinds of dtype specifications
    def __init__(
        self,
        *,
        executors=None,
        devicetypes=None,
        dtypes=None,
        num_devices: int = 1,
        decorators: None | Sequence = None,
        scope=None,
        as_name: str | None = None,
    ):
        self.executors = set(executors) if executors is not None else set(_all_test_executors())
        self.devicetypes = set(devicetypes) if devicetypes is not None else set(available_devicetypes())

        self.devicetypes = set(filter_ci_devicetypes(self.devicetypes))

        if dtypes == NOTHING:
            self.dtypes = (None,)
        else:
            self.dtypes = datatypes.resolve_dtypes(dtypes) if dtypes is not None else datatypes.all_dtypes

        self.num_devices = num_devices

        self.decorators = decorators

        # Acquires the caller's global scope
        if scope is None:
            previous_frame = inspect.currentframe().f_back
            scope = previous_frame.f_globals
        self.scope = scope

        self.as_name = as_name

    # TODO: refactor with the ops class above
    def __call__(self, test_template):
        # NOTE: unlike a typical decorator, this __call__ does not return a function, because it may
        #   (and typically does) instantiate multiple functions from the template it consumes
        #   Since Python doesn't natively support one-to-many function decorators, the produced
        #   functions are directly assigned to the requested scope (the caller's global scope by default)

        for executor, devicetype in product(
            sorted(self.executors, key=lambda x: repr(x)), sorted(self.devicetypes, key=lambda x: repr(x))
        ):
            if executor is None or not executor.is_available():
                continue

            if not executor.supports_devicetype(devicetype):
                continue

            # Identifies devices to run the test on
            available_devices = devices.available_devices()
            filtered_devices = list([x for x in available_devices if x.devicetype is devicetype])
            utils.check(self.num_devices > 0, lambda: f"Received an invalid request for {self.num_devices} devices")

            if devicetype is not devices.DeviceType.CPU and len(filtered_devices) < self.num_devices:
                continue

            device_or_devices = None
            if self.num_devices == 1:
                device_or_devices = devices.Device(devicetype, None)
            else:
                device_or_devices = []
                for idx in range(self.num_devices):
                    dev = devices.Device(devicetype, idx)
                    device_or_devices.append(dev)

            for dtype in sorted(self.dtypes, key=lambda t: repr(t)):
                if dtype is not None and not executor.supports_dtype(dtype):
                    continue

                test = _instantiate_executor_test_template(
                    test_template,
                    self.scope,
                    executor=executor,
                    device_or_devices=device_or_devices,
                    dtype=dtype,
                    as_name=self.as_name,
                )

                # Applies decorators
                if self.decorators is not None:
                    for dec in self.decorators:
                        test = dec(test)

                # Adds the instantiated test to the requested scope
                self.scope[test.__name__] = test


def run_snippet(snippet, opinfo, devicetype, dtype, *args, **kwargs):
    try:
        snippet(*args, **kwargs)
    except Exception as ex:
        exc_info = sys.exc_info()

        # Raises exceptions that occur with pytest, and returns debug information when
        # called otherwise
        # NOTE: PYTEST_CURRENT_TEST is set by pytest
        if "PYTEST_CURRENT_TEST" in os.environ:
            raise ex
        return ex, exc_info, snippet, opinfo, devicetype, dtype, args, kwargs

    return None


def requiresTriton(fn):
    @wraps(fn)
    def _fn(*args, **kwargs):
        if not TRITON_AVAILABLE:
            pytest.skip("Requires Triton")
        return fn(*args, **kwargs)

    return _fn


def requiresCUDA(fn):
    import torch

    @wraps(fn)
    def _fn(*args, **kwargs):
        if not torch.cuda.is_available():
            pytest.skip("Requires CUDA")
        return fn(*args, **kwargs)

    return _fn


def requiresNVFuser(fn):
    @wraps(fn)
    def _fn(*args, **kwargs):
        if not NVFUSER_AVAILABLE:
            pytest.skip("Requires nvFuser")
        return fn(*args, **kwargs)

    return _fn


# A dummy decorator that passes the comparator metadata
class custom_comparator:
    def __init__(self, comparator):
        self.comparator = comparator

    def __call__(self, test_template):
        return test_template


@contextlib.contextmanager
def set_default_dtype_ctx(dtype):
    saved_dtype = torch.get_default_dtype()
    torch.set_default_dtype(dtype)
    try:
        yield
    finally:
        torch.set_default_dtype(saved_dtype)


def has_enough_device_memory(required_memory_bytes: int, cuda_device_id: int = 0) -> bool:
    """
    Check if the specified CUDA device has sufficient total memory for a given requirement.

    Args:
        required_memory_bytes: Amount of memory needed in bytes
        cuda_device_id (int, optional): The ID of the CUDA device to check. Defaults to 0.

    Returns:
        bool: True if the device has sufficient total memory, False otherwise
    """
    if not torch.cuda.is_available():
        return False

    # Get total available memory
    total_memory = torch.cuda.get_device_properties(cuda_device_id).total_memory

    return total_memory > required_memory_bytes


def requiresDeviceMemory(required_memory_bytes: int, cuda_device_id: int = 0):
    def decorator(fn):
        @wraps(fn)
        def _fn(*args, **kwargs):
            if not has_enough_device_memory(required_memory_bytes, cuda_device_id):
                pytest.skip(
                    f"Requires {required_memory_bytes} bytes of memory on device {cuda_device_id} but only {torch.cuda.get_device_properties(cuda_device_id).total_memory} bytes are available"
                )
            return fn(*args, **kwargs)

        return _fn

    return decorator
