"""
This should be imported before entrypoints to correctly configure global options prior to importing packages like torch and cv2.

Use this instead of cli_args to import the args:

>>> from comfy.cmd.main_pre import args

It will enable command line argument parsing. If this isn't desired, you must author your own implementation of these fixes.
"""
import ctypes
import importlib.util
import logging
import os
import shutil
import warnings

# must be set before any transitive import of requests
warnings.filterwarnings("ignore", message=".*doesn't match a supported version")

from .. import options
from ..app import logger
from ..cli_args_types import Configuration
from ..component_model.cuda_env import ensure_pytorch_cuda_alloc_conf, should_skip_cuda_alloc_conf_for_xpu

os.environ['TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL'] = '1'
os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1"
os.environ.setdefault("TORCHINDUCTOR_FX_GRAPH_CACHE", "1")
os.environ.setdefault("TORCHINDUCTOR_AUTOGRAD_CACHE", "1")
os.environ["BITSANDBYTES_NOWELCOME"] = "1"
os.environ["NO_ALBUMENTATIONS_UPDATE"] = "1"
os.environ['HF_HUB_DISABLE_TELEMETRY'] = '1'
os.environ['DO_NOT_TRACK'] = '1'
from ..component_model.torch_cache import setup_torch_compile_cache_dirs
setup_torch_compile_cache_dirs()
if os.name == "nt":
    os.environ['MIMALLOC_PURGE_DELAY'] = '0'

this_logger = logging.getLogger(__name__)

options.enable_args_parsing()
if os.name == "nt":
    logging.getLogger("xformers").addFilter(lambda record: 'A matching Triton is not available' not in record.getMessage())
warnings.filterwarnings("ignore", message="torch.utils._pytree._register_pytree_node is deprecated. Please use torch.utils._pytree.register_pytree_node instead.")
warnings.filterwarnings("ignore", message="Torch was not compiled with flash attention.")
warnings.filterwarnings("ignore", message=".*Torch was not compiled with flash attention.*")
warnings.filterwarnings('ignore', category=FutureWarning, message=r'`torch\.cuda\.amp\.custom_fwd.*')
warnings.filterwarnings("ignore", category=UserWarning, message="Please use the new API settings to control TF32 behavior.*")
warnings.filterwarnings("ignore", message="Importing from timm.models.registry is deprecated, please import via timm.models", category=FutureWarning)
warnings.filterwarnings("ignore", message="Importing from timm.models.layers is deprecated, please import via timm.layers", category=FutureWarning)
warnings.filterwarnings("ignore", message="Inheritance class _InstrumentedApplication from web.Application is discouraged", category=DeprecationWarning)
warnings.filterwarnings("ignore", message="Please import `gaussian_filter` from the `scipy.ndimage` namespace; the `scipy.ndimage.filters` namespace is deprecated", category=DeprecationWarning)
warnings.filterwarnings("ignore", message="The installed version of bitsandbytes was compiled without GPU support")
warnings.filterwarnings("ignore", message=r"The pynvml package is deprecated\..*", category=FutureWarning)

# AMD ROCm: when the local GPU's gfx arch isn't in torch's compiled list but a
# same-family fallback is (e.g. RX 7600/gfx1102 on a wheel built for
# gfx1100/1101 only), set HSA_OVERRIDE_GFX_VERSION before torch loads its HIP
# runtime. Idempotent and skipped if the user already exported the variable.
try:
    from ..component_model.amd_compat import maybe_set_hsa_override as _maybe_set_hsa_override
    _hsa_value = _maybe_set_hsa_override()
    if _hsa_value:
        this_logger.info(
            "Set HSA_OVERRIDE_GFX_VERSION=%s (local AMD GPU arch not in torch's compiled list)",
            _hsa_value,
        )
except Exception:  # noqa: BLE001
    pass
warnings.filterwarnings("ignore", category=UserWarning, message="Unsupported Windows version .* ONNX Runtime supports Windows 10 and above, only.")

# tokenizers >=0.23.0rc0 dropped the `sep=` and `cls=` kwargs from
# RobertaProcessing.__new__ (still positional-only) but transformers
# tokenization_clip.py:117 calls it with those kwargs, breaking CLIP tokenizer
# load. Restore the kwarg form by wrapping the class with a shim.
try:
    from tokenizers import processors as _tk_processors
    _orig_roberta_processing = _tk_processors.RobertaProcessing
    try:
        _tk_processors.RobertaProcessing(sep=('</s>', 2), cls=('<s>', 0))
    except TypeError:
        def _roberta_processing_kwarg_shim(*args, **kwargs):
            if not args and 'sep' in kwargs and 'cls' in kwargs:
                args = (kwargs.pop('sep'), kwargs.pop('cls'))
            return _orig_roberta_processing(*args, **kwargs)
        _tk_processors.RobertaProcessing = _roberta_processing_kwarg_shim
except ImportError:
    pass

log_msg_to_filter = "NOTE: Redirects are currently not supported in Windows or MacOs."
logging.getLogger("torch.distributed.elastic.multiprocessing.redirects").addFilter(
    lambda record: log_msg_to_filter not in record.getMessage()
)
logging.getLogger("alembic.runtime.migration").setLevel(logging.WARNING)
logging.getLogger("alembic").setLevel(logging.WARNING)
logging.getLogger(__name__).addFilter(lambda record: "setup plugin" not in record.getMessage())
logging.getLogger("asyncio").addFilter(lambda record: 'Using selector:' not in record.getMessage())
logging.getLogger("requests_cache").setLevel(logging.ERROR)
logging.getLogger("fsspec").setLevel(logging.WARNING)
logging.getLogger("comfy_kitchen.dispatch").setLevel(logging.WARNING)
logging.getLogger("jax").setLevel(logging.WARNING)
logging.getLogger("aiohttp.access").setLevel(logging.WARNING)

from ..cli_args import args

if args.torch_device is not None:
    # --torch-device takes priority over --cuda-device and --default-device
    _td = args.torch_device
    if _td.startswith("cuda:"):
        _cuda_idx = _td.split(":", 1)[1]
        os.environ['CUDA_VISIBLE_DEVICES'] = _cuda_idx
        os.environ['HIP_VISIBLE_DEVICES'] = _cuda_idx
        os.environ["ASCEND_RT_VISIBLE_DEVICES"] = _cuda_idx
    elif _td.startswith("xpu:"):
        os.environ['ONEAPI_DEVICE_SELECTOR'] = f"level_zero:{_td.split(':', 1)[1]}"
    elif _td.startswith("npu:"):
        os.environ["ASCEND_RT_VISIBLE_DEVICES"] = _td.split(":", 1)[1]
    this_logger.info("Set torch device to: %s", _td)
else:
    if args.default_device is not None:
        default_dev = args.default_device
        devices = list(range(32))
        devices.remove(default_dev)
        devices.insert(0, default_dev)
        devices = ','.join(map(str, devices))
        os.environ['CUDA_VISIBLE_DEVICES'] = str(devices)
        os.environ['HIP_VISIBLE_DEVICES'] = str(devices)

    if args.cuda_device is not None:
        os.environ['CUDA_VISIBLE_DEVICES'] = str(args.cuda_device)
        os.environ['HIP_VISIBLE_DEVICES'] = str(args.cuda_device)
        os.environ["ASCEND_RT_VISIBLE_DEVICES"] = str(args.cuda_device)
        this_logger.info("Set cuda device to: {}".format(args.cuda_device))

if args.deterministic:
    if 'CUBLAS_WORKSPACE_CONFIG' not in os.environ:
        os.environ['CUBLAS_WORKSPACE_CONFIG'] = ":4096:8"

if args.oneapi_device_selector is not None:
    os.environ['ONEAPI_DEVICE_SELECTOR'] = args.oneapi_device_selector
    this_logger.info("Set oneapi device selector to: {}".format(args.oneapi_device_selector))

ensure_pytorch_cuda_alloc_conf(skip_for_xpu=should_skip_cuda_alloc_conf_for_xpu())

try:
    pass
except Exception:
    pass


def _fix_pytorch_240():
    """Fixes pytorch 2.4.0"""
    torch_spec = importlib.util.find_spec("torch")
    for folder in torch_spec.submodule_search_locations:
        lib_folder = os.path.join(folder, "lib")
        test_file = os.path.join(lib_folder, "fbgemm.dll")
        dest = os.path.join(lib_folder, "libomp140.x86_64.dll")
        if os.path.exists(dest):
            break

        try:
            with open(test_file, 'rb') as f:
                contents = f.read()
                # todo: dubious
                if b"libomp140.x86_64.dll" not in contents:
                    break
            try:
                _ = ctypes.cdll.LoadLibrary(test_file)
            except FileNotFoundError:
                this_logger.warning("Detected pytorch version with libomp issue, trying to patch")
                try:
                    shutil.copyfile(os.path.join(lib_folder, "libiomp5md.dll"), dest)
                except Exception as exc_info:
                    this_logger.error("While trying to patch a fix for torch 2.4.0, an error occurred, which means this is unlikely to work", exc_info=exc_info)
        except:
            pass


def _create_tracer():
    from ..component_model import setup as _setup_mod
    _setup_mod.setup_tracing(args)
    from opentelemetry import trace
    return trace.get_tracer(args.otel_service_name)


def _configure_logging():
    logging_level = args.logging_level
    # DETAIL remains available as a console level, but upstream no longer
    # creates a detail log file unless one is explicitly configured.
    logger.setup_logger(logging_level, file_outputs=[])

args: Configuration

_configure_logging()
_fix_pytorch_240()


def _stub_soundfile():
    try:
        import soundfile  # noqa: F401  # noqa: F401
    except (OSError, ImportError):
        import sys
        import types
        sf = types.ModuleType("soundfile")
        sf.__version__ = "0.0.0"
        sys.modules["soundfile"] = sf


_stub_soundfile()

# Patch torchvision InterpolationMode for compatibility with newer torchvision
# builds where InterpolationMode moved. Must happen before third-party imports
# (transformers, spandrel, diffusers) that use the old location.
from .. import torchvision_compat  # noqa: F401

# tracer is a lazy module property: created on first access so that when the
# CLI path (setup_post_torch → setup_tracing) has already initialised the
# global TracerProvider we just hand back a tracer from it instead of trying
# to set_tracer_provider a second time.
from ..component_model.module_property import create_module_properties

_module_properties = create_module_properties()


@_module_properties.getter
def _tracer():
    return _create_tracer()


tracer: ...
