import contextlib
import ctypes
import logging
import os
import sys
import warnings
from pathlib import Path


def _compute_version() -> str:
    """Return the package version without being polluted by unrelated installs.

    Preference order:
    1) If running from a source checkout (VERSION file present at repo root),
       use the dynamic version from version_provider (falls back to plain VERSION).
    2) Otherwise, use importlib.metadata for the installed distribution.
    3) As a last resort, return a dev sentinel.
    """
    try:
        repo_root = Path(__file__).resolve().parent.parent
        version_file = repo_root / "VERSION"
        if version_file.is_file():
            try:
                from version_provider import dynamic_metadata  # type: ignore

                return dynamic_metadata("version")
            except Exception:
                # Fall back to the raw VERSION file if provider isn't available.
                return version_file.read_text().strip()
    except Exception:
        # If any of the above fails, fall through to installed metadata.
        pass

    try:
        from importlib.metadata import version as _dist_version  # py3.8+

        return _dist_version("tilelang")
    except Exception as exc:
        warnings.warn(
            f"tilelang version metadata unavailable ({exc!r}); using development version.",
            RuntimeWarning,
            stacklevel=2,
        )
        return "0.0.dev0"


__version__ = _compute_version()
del _compute_version


logger = logging.getLogger(__name__)


def set_log_level(level):
    """Set the logging level for the module's logger.

    Args:
        level (str or int): Can be the string name of the level (e.g., 'INFO') or the actual level (e.g., logging.INFO).
        OPTIONS: 'DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL'
    """
    if isinstance(level, str):
        level = getattr(logging, level.upper(), logging.INFO)
    logger.setLevel(level)


def _init_logger():
    """Initialize the logger specific for this module with custom settings and a Tqdm-based handler."""
    try:
        from tqdm.auto import tqdm
    except ImportError:
        tqdm = None

    class TqdmLoggingHandler(logging.Handler):
        """Custom logging handler that directs log output to tqdm progress bar to avoid interference."""

        def __init__(self, level=logging.NOTSET):
            """Initialize the handler with an optional log level."""
            super().__init__(level)

        def emit(self, record):
            """Emit a log record. Messages are written to tqdm to ensure output in progress bars isn't corrupted."""
            try:
                msg = self.format(record)
                if tqdm is not None:
                    tqdm.write(msg)
            except Exception:
                self.handleError(record)

    handler = TqdmLoggingHandler()
    formatter = logging.Formatter(
        fmt="%(asctime)s  [TileLang:%(name)s:%(levelname)s] (%(filename)s:%(lineno)d): %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
    )
    handler.setFormatter(formatter)
    logger.addHandler(handler)
    logger.propagate = False
    set_log_level("INFO")


from .env import env as env  # noqa: F401

# Skip logger initialization in light import mode
if not env.is_light_import():
    _init_logger()

del _init_logger


def _disable_tvm_ffi_torch_c_dlpack(torch_module):
    is_rocm = getattr(torch_module.version, "hip", None) is not None
    mps_backend = getattr(getattr(torch_module, "backends", None), "mps", None)
    is_mps = sys.platform == "darwin" and mps_backend is not None and mps_backend.is_available()
    if not is_rocm and not is_mps:
        return

    os.environ.setdefault("TVM_FFI_DISABLE_TORCH_C_DLPACK", "1")
    if is_mps:
        # PyTorch's native exchange API queries a stream handle that MPS cannot provide.
        os.environ.setdefault("TVM_FFI_SKIP_DLPACK_C_EXCHANGE_API", "1")
    try:
        from tvm_ffi import _optional_torch_c_dlpack
    except Exception:
        return

    # TVM Relax calls this loader directly while importing, bypassing the
    # tvm-ffi module-level env guard.
    def _disabled_torch_c_dlpack_extension(*_args, **_kwargs):
        return None

    _optional_torch_c_dlpack.load_torch_c_dlpack_extension = _disabled_torch_c_dlpack_extension


@contextlib.contextmanager
def _lazy_load_lib():
    import torch  # preload torch to avoid dlopen errors

    _disable_tvm_ffi_torch_c_dlpack(torch)

    if sys.platform.startswith("win32"):
        yield
        return

    old_flags = sys.getdlopenflags()
    old_init = ctypes.CDLL.__init__

    def lazy_init(self, name, mode=ctypes.DEFAULT_MODE, *args, **kwargs):
        return old_init(self, name, mode | os.RTLD_LAZY, *args, **kwargs)

    sys.setdlopenflags(old_flags | os.RTLD_LAZY)
    ctypes.CDLL.__init__ = lazy_init
    try:
        yield
    finally:
        sys.setdlopenflags(old_flags)
        ctypes.CDLL.__init__ = old_init


# Skip heavy imports in light import mode
if not env.is_light_import():
    with _lazy_load_lib():
        from .env import (  # noqa: F401
            enable_cache,
            disable_cache,
            is_cache_enabled,
            get_windows_runtime_dll_dirs,
            prepend_dll_search_path,
        )
        from . import libinfo

        if sys.platform.startswith("win32"):
            # Make sibling-package DLLs (tvm_ffi, z3) discoverable via PATH,
            # then register all native dependency dirs with the secure DLL
            # loader used by Python 3.8+ for absolute-path DLL loads.
            runtime_dll_dirs = get_windows_runtime_dll_dirs()
            prepend_dll_search_path(runtime_dll_dirs)
            dll_dirs = dict.fromkeys([*libinfo.get_dll_directories(), *runtime_dll_dirs])
            _dll_handles = [os.add_dll_directory(p) for p in dll_dirs]

        import tvm
        import tvm.base  # noqa: F401
        from tvm import DataType  # noqa: F401

        def _load_tile_lang_lib():
            """Load Tile Lang lib"""
            lib_path = libinfo.find_lib_path("tilelang")
            return ctypes.CDLL(lib_path), lib_path

        # only load once here
        if env.SKIP_LOADING_TILELANG_SO == "0":
            _LIB, _LIB_PATH = _load_tile_lang_lib()

    from .jit import jit, JITKernel, compile, par_compile  # noqa: F401
    from .profiler import Profiler  # noqa: F401
    from .utils import (
        TensorSupplyType,  # noqa: F401
        deprecated,  # noqa: F401
        build_date,  # noqa: F401
    )
    from .layout import (
        Layout,  # noqa: F401
        Fragment,  # noqa: F401
    )
    from . import (
        analysis,  # noqa: F401
        transform,  # noqa: F401
        language,  # noqa: F401
        engine,  # noqa: F401
        tools,  # noqa: F401
    )
    from .language import dtypes  # noqa: F401
    from .autotuner import autotune  # noqa: F401
    from .transform import PassConfigKey  # noqa: F401
    from .engine import lower, register_cuda_postproc, register_hip_postproc, register_c_postproc  # noqa: F401
    from .math import *  # noqa: F403
    from . import ir  # noqa: F401
    from . import tileop  # noqa: F401
    from . import cpu as cpu  # noqa: F401
    from . import cuda as cuda  # noqa: F401
    from . import rocm as rocm  # noqa: F401
    from . import metal as metal  # noqa: F401

del _lazy_load_lib

# Install pass diff hook if TILELANG_PASS_DIFF is enabled (zero overhead when off).
# AutoDD uses a light package import so its CLI can run without loading TVM.
if not env.is_light_import():
    from .utils.pass_diff_hook import install_pass_diff_hook as _install_pass_diff_hook  # noqa: E402

    _install_pass_diff_hook()
    del _install_pass_diff_hook
