import importlib
import inspect
import os

import torch

from optimum.quanto import QModuleMixin, register_qmodule
from optimum.quanto.tensor.qtype import qtype as _quanto_qtype

from . import safetensors2


_QTYPE_QMODULE_CACHE = None
_QMODULE_BASE_ATTRS = None

_DEFAULT_KIND_PRIORITIES = {
    "nvfp4": 1,
    "nunchaku": 2,
    "nunchaku_int4": 2,
    "nunchaku_fp4": 2,
    "fp8": 10,
    "scaled_fp8": 10,
    "scaled_float8_e4m3fn": 10,
    "scaled_float8_e5m2": 10,
    "qfloat8": 10,
    "qfloat8_e4m3fn": 10,
    "qfloat8_e5m2": 10,
    "float8_e4m3fn": 10,
    "float8_e5m2": 10,
    "int8": 11,
    "qint8": 11,
}


def _extract_qtypes(handler):
    for obj in vars(handler).values():
        if isinstance(obj, _quanto_qtype):
            yield obj


def _extract_qmodule_classes(handler):
    for obj in vars(handler).values():
        if inspect.isclass(obj) and issubclass(obj, QModuleMixin) and issubclass(obj, torch.nn.Linear):
            if obj is QLinearQuantoRouter:
                continue
            yield obj


def _build_qmodule_cache():
    mapping = {}
    for handler in _load_handlers():
        qmodule_classes = list(_extract_qmodule_classes(handler))
        if len(qmodule_classes) != 1:
            continue
        qmodule_cls = qmodule_classes[0]
        for qt in _extract_qtypes(handler):
            mapping.setdefault(qt, qmodule_cls)
    return mapping


def _get_qmodule_base_attrs():
    global _QMODULE_BASE_ATTRS
    if _QMODULE_BASE_ATTRS is not None:
        return _QMODULE_BASE_ATTRS
    base = torch.nn.Linear(1, 1, bias=True)
    _QMODULE_BASE_ATTRS = set(base.__dict__.keys())
    _QMODULE_BASE_ATTRS.update({
        "_parameters",
        "_buffers",
        "_modules",
        "_non_persistent_buffers_set",
    })
    return _QMODULE_BASE_ATTRS


def _get_qmodule_for_qtype(qtype_obj):
    global _QTYPE_QMODULE_CACHE
    if qtype_obj is None:
        return None
    if _QTYPE_QMODULE_CACHE is None or qtype_obj not in _QTYPE_QMODULE_CACHE:
        _QTYPE_QMODULE_CACHE = _build_qmodule_cache()
    return _QTYPE_QMODULE_CACHE.get(qtype_obj)


def _load_with_qmodule(
    module,
    qmodule_cls,
    state_dict,
    prefix,
    local_metadata,
    strict,
    missing_keys,
    unexpected_keys,
    error_msgs,
):
    device = module.weight.device if torch.is_tensor(module.weight) else None
    if torch.is_tensor(module.weight) and module.weight.dtype.is_floating_point:
        weight_dtype = module.weight.dtype
    elif torch.is_tensor(getattr(module, "bias", None)) and module.bias.dtype.is_floating_point:
        weight_dtype = module.bias.dtype
    else:
        weight_dtype = torch.float16
    tmp = qmodule_cls(
        module.in_features,
        module.out_features,
        bias=module.bias is not None,
        device=device,
        dtype=weight_dtype,
        weights=module.weight_qtype,
        activations=module.activation_qtype,
        optimizer=module.optimizer,
        quantize_input=True,
    )
    setter = getattr(tmp, "set_default_dtype", None)
    if callable(setter):
        setter(getattr(module, "_router_default_dtype", None) or module.weight.dtype)
    tmp._load_from_state_dict(
        state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
    )

    module.weight = tmp.weight
    module.bias = tmp.bias
    module.input_scale = tmp.input_scale
    module.output_scale = tmp.output_scale

    ignore = set(_get_qmodule_base_attrs())
    ignore.update({
        "_quantize_hooks",
        "training",
        "_router_default_dtype",
    })
    for name, value in tmp.__dict__.items():
        if name in ignore:
            continue
        setattr(module, name, value)
    module._router_forward_impl = qmodule_cls.forward


@register_qmodule(torch.nn.Linear)
class QLinearQuantoRouter(QModuleMixin, torch.nn.Linear):
    @classmethod
    def qcreate(
        cls,
        module,
        weights,
        activations=None,
        optimizer=None,
        device=None,
    ):
        if torch.is_tensor(module.weight) and module.weight.dtype.is_floating_point:
            weight_dtype = module.weight.dtype
        elif torch.is_tensor(getattr(module, "bias", None)) and module.bias.dtype.is_floating_point:
            weight_dtype = module.bias.dtype
        else:
            weight_dtype = torch.float16
        return cls(
            module.in_features,
            module.out_features,
            module.bias is not None,
            device=device,
            dtype=weight_dtype,
            weights=weights,
            activations=activations,
            optimizer=optimizer,
            quantize_input=True,
        )

    def __init__(
        self,
        in_features,
        out_features,
        bias=True,
        device=None,
        dtype=None,
        weights=None,
        activations=None,
        optimizer=None,
        quantize_input=True,
    ):
        super().__init__(
            in_features,
            out_features,
            bias=bias,
            device=device,
            dtype=dtype,
            weights=weights,
            activations=activations,
            optimizer=optimizer,
            quantize_input=quantize_input,
        )
        self._router_default_dtype = dtype

    def set_default_dtype(self, dtype):
        self._router_default_dtype = dtype

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        impl = getattr(self, "_router_forward_impl", None)
        if impl is not None:
            return impl(self, input)
        return torch.nn.functional.linear(input, self.qweight, bias=self.bias)

    def _load_from_state_dict(
        self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
    ):
        qmodule_cls = _get_qmodule_for_qtype(self.weight_qtype)
        if qmodule_cls is not None and qmodule_cls is not QLinearQuantoRouter:
            return _load_with_qmodule(
                self, qmodule_cls, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
            )
        return super()._load_from_state_dict(
            state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
        )


_FP8_QUANTO_BRIDGE_MODULE = ".fp8_quanto_bridge"

_HANDLER_MODULES = [
    _FP8_QUANTO_BRIDGE_MODULE,
]
_HANDLER_OBJECTS = []


def register_handler(handler):
    global _QTYPE_QMODULE_CACHE
    if isinstance(handler, str):
        if handler not in _HANDLER_MODULES:
            _HANDLER_MODULES.append(handler)
            _QTYPE_QMODULE_CACHE = None
        return handler
    if handler not in _HANDLER_OBJECTS:
        _HANDLER_OBJECTS.append(handler)
        _QTYPE_QMODULE_CACHE = None
    return handler


def unregister_handler(handler):
    global _QTYPE_QMODULE_CACHE
    removed = False
    if isinstance(handler, str):
        if handler in _HANDLER_MODULES:
            _HANDLER_MODULES.remove(handler)
            removed = True
    elif handler in _HANDLER_OBJECTS:
        _HANDLER_OBJECTS.remove(handler)
        removed = True
    if removed:
        _QTYPE_QMODULE_CACHE = None
    return removed


def _load_handlers():
    handlers = []
    for mod_path in _HANDLER_MODULES:
        module = importlib.import_module(mod_path, package=__package__)
        if not hasattr(module, "detect") or not hasattr(module, "convert_to_quanto"):
            raise RuntimeError(
                f"Quant handler '{mod_path}' must define detect() and convert_to_quanto() functions."
            )
        handlers.append(module)
    for handler in _HANDLER_OBJECTS:
        if not hasattr(handler, "detect") or not hasattr(handler, "convert_to_quanto"):
            raise RuntimeError(
                "Quant handler object must define detect() and convert_to_quanto() functions."
            )
        handlers.append(handler)
    register_qmodule(torch.nn.Linear)(QLinearQuantoRouter)
    return handlers


def _handler_name(handler):
    return getattr(handler, "HANDLER_NAME", handler.__name__.split(".")[-1])

_FILE_EXTENSION_HANDLERS = {}


def register_file_extension(extension, handler):
    if not extension or handler is None:
        return
    ext = str(extension).lower()
    if ext.startswith("."):
        ext = ext[1:]
    if not ext:
        return
    _FILE_EXTENSION_HANDLERS[ext] = handler


def get_extension_handler(file_path):
    if not isinstance(file_path, str):
        return None
    ext = os.path.splitext(file_path)[1].lower().lstrip(".")
    if not ext:
        return None
    return _FILE_EXTENSION_HANDLERS.get(ext)


def normalize_extension_path(file_path):
    handler = get_extension_handler(file_path)
    if handler is None:
        return file_path
    normalizer = getattr(handler, "normalize", None)
    if not callable(normalizer):
        ext = os.path.splitext(file_path)[1].lower().lstrip(".")
        raise Exception(f"Missing normalize for *.{ext} handler")
    return normalizer(file_path)

def _normalize_kind_key(value):
    if value is None:
        return ""
    if isinstance(value, _quanto_qtype):
        return value.name.lower()
    name = getattr(value, "name", None)
    if isinstance(name, str) and name:
        return name.lower()
    return str(value).lower()


def _priority_for_kind(kind):
    key = _normalize_kind_key(kind)
    if not key:
        return None
    if key in _DEFAULT_KIND_PRIORITIES:
        return _DEFAULT_KIND_PRIORITIES[key]
    if "nunchaku" in key:
        return _DEFAULT_KIND_PRIORITIES["nunchaku"]
    if "float8" in key or "fp8" in key:
        return _DEFAULT_KIND_PRIORITIES["fp8"]
    if "int8" in key:
        return _DEFAULT_KIND_PRIORITIES["int8"]
    return None


def _get_handler_priority(handler):
    for attr in ("HANDLER_PRIORITY", "PRIORITY", "priority"):
        value = getattr(handler, attr, None)
        if isinstance(value, (int, float)):
            return int(value)
    return _priority_for_kind(_handler_name(handler))


def _select_primary_kind(names, priority_map=None):
    if not names:
        return None
    best_name = None
    best_priority = None
    for name in names:
        priority = None
        if priority_map is not None:
            priority = priority_map.get(name)
        if priority is None:
            priority = _priority_for_kind(name)
        if priority is None:
            priority = 1000
        if best_priority is None or priority < best_priority:
            best_priority = priority
            best_name = name
    return best_name or names[0]


def _merge_quant_maps(target, source):
    if not source:
        return target
    if target is None:
        target = {}
    for key, cfg in source.items():
        if key not in target:
            target[key] = cfg
            continue
        if target[key] == cfg:
            continue
        current_priority = _priority_for_kind((target[key] or {}).get("weights")) if isinstance(target[key], dict) else None
        incoming_priority = _priority_for_kind((cfg or {}).get("weights")) if isinstance(cfg, dict) else None
        if current_priority is None:
            current_priority = 1000
        if incoming_priority is None:
            incoming_priority = 1000
        if incoming_priority < current_priority:
            target[key] = cfg
    return target


def _as_field_tuple(value):
    if not value:
        return ()
    if isinstance(value, str):
        return (value,)
    return tuple(value)


def _get_split_handler(info, field, default_handlers):
    handlers = info.get("split_handlers") or info.get("field_handlers") or {}
    if handlers:
        handler = handlers.get(field)
        if handler is not None:
            return handler
    if default_handlers:
        return default_handlers.get(field)
    return None


def _get_split_base_fields(info, split_fields):
    base_fields = _as_field_tuple(info.get("base_fields") or info.get("base_field"))
    if base_fields:
        return base_fields
    if split_fields:
        return (next(iter(split_fields.keys())),)
    return ()


def _merge_share_fields(info, share_fields):
    info_fields = _as_field_tuple(info.get("share_fields") or info.get("shared_fields"))
    return tuple(sorted(set(info_fields).union(_as_field_tuple(share_fields))))


def _call_split_handler(handler, *, src, dim, split_sizes, context):
    if handler is None:
        return None
    try:
        chunks = handler(src=src, dim=dim, split_sizes=split_sizes, context=context)
    except Exception:
        return None
    if not isinstance(chunks, (list, tuple)) or len(chunks) != len(split_sizes):
        return None
    return chunks


def _fill_sub_maps(sub_maps, name, value):
    for sub_map in sub_maps:
        sub_map[name] = value


def _get_quantized_subtensors(p):
    getter = getattr(p, "get_quantized_subtensors", None)
    if getter is None:
        return None
    sub_tensors = getter()
    if not sub_tensors:
        return None
    if isinstance(sub_tensors, dict):
        sub_tensors = list(sub_tensors.items())
    out = []
    for name, tensor in sub_tensors:
        if tensor is None:
            continue
        if torch.is_tensor(tensor):
            out.append((name, tensor))
    return out if out else None


def sd_split_linear(
    state_dict,
    split_map,
    split_fields=None,
    share_fields=None,
    verboseLevel=1,
    split_handlers=None,
    allowed_bases=None,
    return_split_bases=False,
):
    if not split_map:
        return (state_dict, []) if return_split_bases else state_dict
    split_fields = split_fields or {}
    share_fields = share_fields or ()
    split_handlers = split_handlers or {}
    base_fields_by_suffix = {
        suffix: _get_split_base_fields(info or {}, split_fields)
        for suffix, info in split_map.items()
    }
    def _skip(msg):
        pass

    bases = {}
    for key in state_dict.keys():
        for suffix, base_fields in base_fields_by_suffix.items():
            for base_field in base_fields:
                suffix_token = f"{suffix}.{base_field}"
                if not key.endswith(suffix_token):
                    continue
                base = key[: -len("." + base_field)]
                if base.endswith(suffix):
                    bases[base] = suffix
                break

    if allowed_bases is not None:
        allowed_set = set(allowed_bases)
        bases = {base: suffix for base, suffix in bases.items() if base in allowed_set}

    if not bases:
        return (state_dict, []) if return_split_bases else state_dict

    split_bases = []

    for base, suffix in bases.items():
        info = split_map.get(suffix) or {}
        mapped = info.get("mapped_modules") or info.get("mapped_suffixes") or info.get("mapped") or []
        if not mapped:
            continue

        base_fields = base_fields_by_suffix.get(suffix) or _get_split_base_fields(info, split_fields)
        size_field = info.get("size_field") or (base_fields[0] if base_fields else None)
        size_tensor = state_dict.get(base + "." + size_field) if size_field else None
        split_dim = info.get("split_dim", 0)
        split_sizes = list(info.get("split_sizes") or [])
        if not split_sizes:
            if size_tensor is None:
                continue
            if size_tensor.dim() <= split_dim:
                _skip(f"{base}: dim={size_tensor.dim()} split_dim={split_dim}")
                continue
            out_dim = size_tensor.size(split_dim)
            if out_dim % len(mapped) != 0:
                _skip(f"{base}: out_dim={out_dim} not divisible by {len(mapped)}")
                continue
            split_sizes = [out_dim // len(mapped)] * len(mapped)
        elif None in split_sizes:
            if size_tensor is None:
                continue
            if size_tensor.dim() <= split_dim:
                _skip(f"{base}: dim={size_tensor.dim()} split_dim={split_dim}")
                continue
            known = sum(size for size in split_sizes if size is not None)
            none_count = split_sizes.count(None)
            remaining = size_tensor.size(split_dim) - known
            if remaining < 0 or remaining % none_count != 0:
                _skip(f"{base}: cannot resolve split sizes")
                continue
            fill = remaining // none_count
            split_sizes = [fill if size is None else size for size in split_sizes]

        total = sum(split_sizes)
        prefix = base[: -len(suffix)]
        target_bases = [prefix + name for name in mapped]
        added = 0

        field_tensors = {
            field: state_dict.get(base + "." + field)
            for field in set(split_fields.keys()).union(share_fields)
        }
        base_ctx = {
            "state_dict": state_dict,
            "base": base,
            "suffix": suffix,
            "split_sizes": split_sizes,
            "total": total,
            "mapped": mapped,
            "target_bases": target_bases,
            "verboseLevel": verboseLevel,
            "split_fields": split_fields,
            "share_fields": share_fields,
            "field_tensors": field_tensors,
            "size_field": size_field,
            "size_tensor": size_tensor,
            "split_dim": split_dim,
            "info": info,
        }
        fields_iter = list(split_fields.items()) + [(field, None) for field in share_fields]
        for field, dim in fields_iter:
            src = field_tensors.get(field)
            if src is None:
                continue
            if dim is None:
                for target_base in target_bases:
                    dest_key = target_base + "." + field
                    if dest_key not in state_dict:
                        state_dict[dest_key] = src
                        added += 1
                continue
            handler = _get_split_handler(info, field, split_handlers)
            chunks = _call_split_handler(
                handler,
                src=src,
                dim=dim,
                split_sizes=split_sizes,
                context=dict(base_ctx, field=field),
            )
            if chunks is None:
                if src.dim() <= dim:
                    _skip(f"{base}.{field}: dim={src.dim()} split_dim={dim}")
                    continue
                if src.size(dim) != total:
                    _skip(f"{base}.{field}: size({dim})={src.size(dim)} expected={total}")
                    continue
                chunks = torch.split(src, split_sizes, dim=dim)
            for target_base, chunk in zip(target_bases, chunks):
                if torch.is_tensor(chunk) and not chunk.is_contiguous():
                    chunk = chunk.contiguous()
                dest_key = target_base + "." + field
                if dest_key not in state_dict:
                    state_dict[dest_key] = chunk
                    added += 1

        if added:
            for field in list(split_fields.keys()) + list(share_fields):
                state_dict.pop(base + "." + field, None)
            split_bases.append(base)
    if return_split_bases:
        return state_dict, split_bases
    return state_dict


def split_linear_modules(model, map, split_handlers=None, share_fields=None):
    from accelerate import init_empty_weights

    split_handlers = split_handlers or {}
    share_fields = share_fields or ()

    modules_dict = { k: m for k, m in model.named_modules()}
    for module_suffix, split_info in map.items():
        mapped_modules = split_info["mapped_modules"]
        split_sizes = split_info["split_sizes"]
        split_share_fields = _merge_share_fields(split_info, share_fields)
        split_dims = split_info.get("split_dims") or {}
        for k, module in modules_dict.items():
            if k.endswith("." + module_suffix):
                parent_module = modules_dict[k[:len(k)-len(module_suffix)-1]]
                weight = module.weight
                bias = getattr(module, "bias", None)
                if isinstance(module, QModuleMixin):
                    out_features_total = weight.size(0)
                    if sum(split_sizes) != out_features_total:
                        raise ValueError(
                            f"Split sizes {split_sizes} do not match out_features {out_features_total} for '{k}'."
                        )
                    in_features = weight.size(1)
                    if bias is not None and bias.dim() > 0 and bias.size(0) == out_features_total:
                        sub_biases = torch.split(bias, split_sizes, dim=0)
                    else:
                        sub_biases = [bias] * len(split_sizes)

                    sub_tensors = _get_quantized_subtensors(weight)
                    if not sub_tensors:
                        raise ValueError(f"Unable to split quantized weight for '{k}'.")
                    sub_maps = [dict() for _ in split_sizes]
                    field_tensors = {name: tensor for name, tensor in sub_tensors}
                    base_ctx = {
                        "module": module,
                        "module_name": k,
                        "module_suffix": module_suffix,
                        "mapped_modules": mapped_modules,
                        "split_sizes": split_sizes,
                        "out_features": out_features_total,
                        "in_features": in_features,
                        "field_tensors": field_tensors,
                        "info": split_info,
                    }
                    for name, tensor in sub_tensors:
                        if tensor is None or name in split_share_fields or tensor.dim() <= 1:
                            _fill_sub_maps(sub_maps, name, tensor)
                            continue
                        split_dim = split_dims.get(name)
                        if split_dim is None:
                            if tensor.size(0) == out_features_total:
                                split_dim = 0
                            elif tensor.dim() > 1 and tensor.size(1) == out_features_total:
                                split_dim = 1
                            else:
                                split_dim = 0
                        handler = _get_split_handler(split_info, name, split_handlers)
                        chunks = _call_split_handler(
                            handler,
                            src=tensor,
                            dim=split_dim,
                            split_sizes=split_sizes,
                            context=dict(base_ctx, split_dim=split_dim),
                        )
                        if chunks is None:
                            if tensor.dim() <= split_dim or tensor.size(split_dim) != out_features_total:
                                got_size = "n/a" if tensor.dim() <= split_dim else tensor.size(split_dim)
                                raise ValueError(
                                    f"Cannot split '{k}' quantized tensor '{name}': "
                                    f"expected size({split_dim})={out_features_total}, got {got_size}."
                                )
                            chunks = torch.split(tensor, split_sizes, dim=split_dim)
                        for sub_map, chunk in zip(sub_maps, chunks):
                            sub_map[name] = chunk

                    create_fn = getattr(weight.__class__, "create", None)
                    if not callable(create_fn):
                        raise ValueError(f"Quantized weight class '{weight.__class__.__name__}' has no create()")
                    create_sig = inspect.signature(create_fn)
                    base_kwargs = {
                        "qtype": getattr(weight, "qtype", None),
                        "axis": getattr(weight, "axis", None),
                        "stride": weight.stride(),
                        "dtype": weight.dtype,
                        "activation_qtype": getattr(weight, "activation_qtype", None),
                        "requires_grad": weight.requires_grad,
                        "group_size": getattr(weight, "_group_size", None),
                        "device": weight.device,
                    }

                    qmodule_cls = module.__class__
                    for sub_name, sub_size, sub_map, sub_bias in zip(
                        mapped_modules, split_sizes, sub_maps, sub_biases
                    ):
                        with init_empty_weights():
                            sub_module = qmodule_cls(
                                in_features,
                                sub_size,
                                bias=bias is not None,
                                device="cpu",
                                dtype=weight.dtype,
                                weights=module.weight_qtype,
                                activations=module.activation_qtype,
                                optimizer=module.optimizer,
                                quantize_input=True,
                            )
                        size = list(weight.size())
                        if size:
                            size[0] = sub_size
                        base_kwargs["size"] = tuple(size)
                        create_kwargs = {}
                        missing = []
                        for name, param in create_sig.parameters.items():
                            if name == "self":
                                continue
                            if name in sub_map:
                                create_kwargs[name] = sub_map[name]
                            elif name in base_kwargs and base_kwargs[name] is not None:
                                create_kwargs[name] = base_kwargs[name]
                            elif param.default is param.empty:
                                missing.append(name)
                        if missing:
                            raise ValueError(
                                f"Unable to rebuild quantized weight for '{k}.{sub_name}': "
                                f"missing {missing}."
                            )
                        sub_weight = create_fn(**create_kwargs)
                        sub_module.weight = torch.nn.Parameter(sub_weight, requires_grad=weight.requires_grad)
                        if sub_bias is not None:
                            sub_module.bias = torch.nn.Parameter(sub_bias)
                        sub_module.optimizer = module.optimizer
                        sub_module.weight_qtype = module.weight_qtype
                        sub_module.activation_qtype = module.activation_qtype
                        setattr(parent_module, sub_name, sub_module)
                else:
                    sub_data = torch.split(weight, split_sizes, dim=0)
                    sub_bias = torch.split(bias, split_sizes, dim=0) if bias is not None else [None] * len(split_sizes)
                    for sub_name, sub_weight, sub_biases in zip(mapped_modules, sub_data, sub_bias):
                        sub_module = torch.nn.Linear(
                            module.in_features,
                            sub_weight.size(0),
                            bias=bias is not None,
                            device=weight.device,
                            dtype=weight.dtype,
                        )
                        sub_module.weight = torch.nn.Parameter(sub_weight)
                        if sub_biases is not None:
                            sub_module.bias = torch.nn.Parameter(sub_biases)
                        setattr(parent_module, sub_name, sub_module)
                delattr(parent_module, module_suffix)


def detect_safetensors_format(state_dict, verboseLevel=1):
    matches = []
    details = {}
    priorities = {}
    for handler in _load_handlers():
        result = handler.detect(state_dict, verboseLevel=verboseLevel)
        name = _handler_name(handler)
        details[name] = result
        if result.get("matched", False):
            matches.append(name)
            priorities[name] = _get_handler_priority(handler)
    kind = _select_primary_kind(matches, priorities) or "none"
    return {"kind": kind, "found": matches, "details": details, "mixed": len(matches) > 1}


def detect_and_convert(state_dict, default_dtype, verboseLevel=1):
    info = detect_safetensors_format(state_dict, verboseLevel=verboseLevel)
    kind = info.get("kind", "none")
    matches = info.get("found", []) or []
    if kind in ("none", "quanto") and not matches:
        return {"state_dict": state_dict, "quant_map": {}, "kind": kind, "details": info}

    handlers = _load_handlers()
    handler_map = {_handler_name(handler): handler for handler in handlers}
    if not matches:
        raise RuntimeError(f"Unsupported quantization format '{kind}'")

    if len(matches) == 1:
        handler = handler_map.get(matches[0])
        if handler is None:
            raise RuntimeError(f"Unsupported quantization format '{kind}'")
        detection = info.get("details", {}).get(matches[0], {})
        conv = handler.convert_to_quanto(
            state_dict,
            default_dtype=default_dtype,
            verboseLevel=verboseLevel,
            detection=detection,
        )
        conv["kind"] = kind
        conv["details"] = info
        return conv

    def _match_priority(name):
        handler = handler_map.get(name)
        priority = _get_handler_priority(handler) if handler is not None else None
        if priority is None:
            priority = _priority_for_kind(name)
        if priority is None:
            priority = 1000
        return priority

    ordered_matches = sorted(matches, key=_match_priority)
    merged_state = state_dict
    merged_map = {}
    for name in ordered_matches:
        handler = handler_map.get(name)
        if handler is None:
            continue
        detection = info.get("details", {}).get(name, {})
        conv = handler.convert_to_quanto(
            merged_state,
            default_dtype=default_dtype,
            verboseLevel=verboseLevel,
            detection=detection,
        )
        merged_state = conv.get("state_dict", merged_state)
        merged_map = _merge_quant_maps(merged_map, conv.get("quant_map", {}))
    return {"state_dict": merged_state, "quant_map": merged_map, "kind": kind, "details": info}


def get_available_qtypes():
    try:
        from optimum.quanto.tensor.qtype import qtypes as _quanto_qtypes
    except Exception:
        return []
    return sorted(_quanto_qtypes.keys())


def get_available_qtype_aliases():
    aliases = set()
    for name in get_available_qtypes():
        key = str(name).lower()
        aliases.add(key)
        if key.startswith("q") and len(key) > 1:
            aliases.add(key[1:])
        if "float8" in key:
            aliases.add("fp8")
    return aliases


def get_quantization_tokens(quantization):
    if quantization is None:
        return []
    key = str(quantization).lower()
    if len(key) == 0:
        return []
    aliases = get_available_qtype_aliases()
    if key not in aliases:
        return []
    tokens = {key}
    if key.startswith("q") and len(key) > 1:
        tokens.add(key[1:])
    if "float8" in key or key == "fp8":
        tokens.add("fp8")
    if "int4" in key:
        tokens.add("int4")
    if "int8" in key:
        tokens.add("int8")
    return sorted(tokens, key=len, reverse=True)


def get_quantization_label(quantization):
    if quantization is None:
        return ""
    key = str(quantization).lower()
    if key in ("", "none", "bf16", "fp16", "float16", "bfloat16"):
        return ""
    aliases = get_available_qtype_aliases()
    if key not in aliases:
        return ""
    if "float8" in key or key == "fp8":
        return "FP8"
    if key.startswith("q"):
        key = key[1:]
    return key.replace("_", " ").upper()


_quantization_filename_cache = {}


def _normalize_quant_file_key(file_path):
    try:
        return os.path.normcase(os.path.abspath(file_path))
    except Exception:
        return str(file_path).lower()


def get_cached_quantization_for_file(file_path):
    if not file_path:
        return None
    return _quantization_filename_cache.get(_normalize_quant_file_key(file_path))


def cache_quantization_for_file(file_path, kind):
    if not file_path or not kind:
        return
    key = _normalize_quant_file_key(file_path)
    if key not in _quantization_filename_cache:
        _quantization_filename_cache[key] = kind


def _detect_kind_from_handlers(file_path, verboseLevel=1):
    found = []
    for handler in _load_handlers():
        fn = getattr(handler, "detect_quantization_kind_for_file", None)
        if fn is None:
            continue
        try:
            kind = fn(file_path, verboseLevel=verboseLevel)
        except TypeError:
            kind = fn(file_path)
        if kind:
            found.append((kind, _get_handler_priority(handler)))
    if not found:
        return None
    found.sort(key=lambda entry: entry[1] if entry[1] is not None else 1000)
    return found[0][0]


def _detect_label_from_handlers(file_path, verboseLevel=0):
    for handler in _load_handlers():
        fn = getattr(handler, "detect_quantization_label_from_filename", None)
        if fn is None:
            continue
        try:
            label = fn(file_path, verboseLevel=verboseLevel)
        except TypeError:
            label = fn(file_path)
        if label:
            return label
    return ""


def _infer_qtype_from_quantization_map(quantization_map):
    if not quantization_map:
        return None
    counts = {}
    for entry in quantization_map.values():
        if not isinstance(entry, dict):
            continue
        weights = entry.get("weights")
        if not weights or weights == "none":
            continue
        key = _normalize_kind_key(weights)
        if not key:
            continue
        counts[key] = counts.get(key, 0) + 1
    if not counts:
        return None
    best_key = None
    best_priority = None
    best_count = None
    for key, count in counts.items():
        priority = _priority_for_kind(key)
        if priority is None:
            priority = 1000
        if (
            best_priority is None
            or priority < best_priority
            or (priority == best_priority and (best_count is None or count > best_count))
        ):
            best_priority = priority
            best_count = count
            best_key = key
    return best_key


def load_metadata_state_dict(file_path):

    if isinstance(file_path, str) and get_extension_handler(file_path) is not None:
        ext_handler = get_extension_handler(file_path)
        if hasattr(ext_handler, "get_file_metadata"):
            return ext_handler.get_file_metadata(file_path)
        else:
            return {}, {}
    state_dict, metadata = safetensors2.load_metadata_state_dict(file_path)
    return state_dict, metadata

def detect_quantization_kind_for_file(file_path, verboseLevel=1):
    cached = get_cached_quantization_for_file(file_path)
    if cached:
        return cached
    if not file_path or not os.path.isfile(file_path):
        return None
    if not (".safetensors" in file_path or ".sft" in file_path):
        kind = _detect_kind_from_handlers(file_path, verboseLevel=verboseLevel)
        if kind:
            cache_quantization_for_file(file_path, kind)
            return kind
        return None

    def _load_full():
        state_dict = {}
        with safetensors2.safe_open(
            file_path,
            framework="pt",
            device="cpu",
            writable_tensors=False,
        ) as f:
            for key in f.keys():
                state_dict[key] = f.get_tensor(key)
            metadata = f.metadata()
        return state_dict, metadata

    def _try_detect(state_dict):
        try:
            info = detect_safetensors_format(state_dict, verboseLevel=verboseLevel)
            return info.get("kind"), True
        except Exception:
            return None, False

    metadata_only = False
    try:
        state_dict, metadata = load_metadata_state_dict(file_path)
        metadata_only = True
    except Exception:
        try:
            state_dict, metadata = _load_full()
        except Exception:
            return None

    kind, ok = _try_detect(state_dict)
    if metadata_only and not ok:
        try:
            state_dict, metadata = _load_full()
            kind, ok = _try_detect(state_dict)
        except Exception:
            kind = None

    if (not kind or kind == "none") and metadata is not None:
        inferred = _infer_qtype_from_quantization_map(metadata.get("quantization_map"))
        if inferred:
            kind = inferred

    cache_quantization_for_file(file_path, kind or "none")
    return kind


def detect_quantization_label_from_filename(filename):
    if not filename:
        return ""
    label = _detect_label_from_handlers(filename, verboseLevel=0)
    if label:
        return label
    cached = get_cached_quantization_for_file(filename)
    if cached:
        return get_quantization_label(cached)
    kind = detect_quantization_kind_for_file(filename, verboseLevel=0)
    if kind:
        label = get_quantization_label(kind)
        if label:
            return label
    base = os.path.basename(filename).lower()
    for token in sorted(get_available_qtype_aliases(), key=len, reverse=True):
        if token and token in base:
            return get_quantization_label(token)
    if "quanto" in base:
        return "QUANTO"
    return ""


def _get_qtype_name_from_quant_map(entry):
    if entry is None:
        return ""
    if isinstance(entry, dict):
        entry = entry.get("weights")
    return _normalize_kind_key(entry)


def _build_qtype_handler_map():
    mapping = {}
    for handler in _load_handlers():
        for qt in _extract_qtypes(handler):
            name = getattr(qt, "name", None)
            if isinstance(name, str) and name:
                mapping.setdefault(name.lower(), handler)
    return mapping


def _collect_fused_bases(state_dict, fused_split_map):
    if not fused_split_map or not state_dict:
        return {}
    suffixes = list(fused_split_map.keys())
    if not suffixes:
        return {}
    bases = {}
    for key in state_dict.keys():
        if key.endswith(".weight"):
            base = key[:-7]
        elif key.endswith(".qweight"):
            base = key[:-8]
        else:
            continue
        for suffix in suffixes:
            if base.endswith(suffix):
                bases[base] = suffix
                break
    return bases


def _remap_quantization_entries(quantization_map, base, targets):
    if not quantization_map or not targets:
        return
    for suffix in ("", ".weight"):
        key = base + suffix if suffix else base
        if key not in quantization_map:
            continue
        cfg = quantization_map.pop(key)
        for target in targets:
            quantization_map[target + suffix] = cfg


def split_fused_weights(state_dict, quantization_map, fused_split_map, default_dtype=None, verboseLevel=1):
    if not fused_split_map or not state_dict:
        return state_dict, quantization_map

    quantization_map = quantization_map or {}
    fused_bases = _collect_fused_bases(state_dict, fused_split_map)
    if not fused_bases:
        return state_dict, quantization_map

    handler_map = _build_qtype_handler_map()
    bases_info = {}
    for base, suffix in fused_bases.items():
        info = fused_split_map.get(suffix) or {}
        mapped = info.get("mapped_modules") or info.get("mapped_suffixes") or info.get("mapped") or []
        if not mapped:
            continue
        prefix = base[:-len(suffix)]
        targets = [prefix + name for name in mapped]
        qtype_name = _get_qtype_name_from_quant_map(quantization_map.get(base))
        if not qtype_name:
            qtype_name = _get_qtype_name_from_quant_map(quantization_map.get(base + ".weight"))
        if not qtype_name:
            weight = state_dict.get(base + ".weight")
            if torch.is_tensor(weight) and getattr(weight, "tensor_type", None) is not None:
                qtype_name = "gguf"
        bases_info[base] = {
            "suffix": suffix,
            "targets": targets,
            "qtype": qtype_name,
        }

    if not bases_info:
        return state_dict, quantization_map

    handler_bases = {}
    for base, info in bases_info.items():
        qtype_name = info["qtype"]
        if not qtype_name:
            continue
        handler = handler_map.get(qtype_name)
        if handler is None:
            continue
        if not callable(getattr(handler, "split_fused_weights", None)):
            continue
        handler_bases.setdefault(handler, set()).add(base)

    handled_bases = set()
    for handler, bases in handler_bases.items():
        fn = getattr(handler, "split_fused_weights", None)
        if not callable(fn):
            continue
        result = fn(
            state_dict,
            fused_split_map,
            quantization_map=quantization_map,
            allowed_bases=bases,
            default_dtype=default_dtype,
            verboseLevel=verboseLevel,
        )
        if isinstance(result, tuple) and len(result) == 2:
            state_dict, split_bases = result
        else:
            state_dict = result or state_dict
            split_bases = []
        handled_bases.update(split_bases or [])

    default_bases = [
        base for base, info in bases_info.items()
        if not info["qtype"] and base not in handled_bases
    ]
    if default_bases:
        state_dict, split_bases = sd_split_linear(
            state_dict,
            fused_split_map,
            split_fields={"weight": 0, "bias": 0},
            verboseLevel=verboseLevel,
            allowed_bases=default_bases,
            return_split_bases=True,
        )
        handled_bases.update(split_bases or [])

    if quantization_map and handled_bases:
        for base in handled_bases:
            info = bases_info.get(base)
            if not info:
                continue
            _remap_quantization_entries(quantization_map, base, info["targets"])

    return state_dict, quantization_map


def apply_pre_quantization(model, state_dict, quantization_map, default_dtype=None, verboseLevel=1):
    remaining = dict(quantization_map or {})
    post_load = []
    for handler in _load_handlers():
        fn = getattr(handler, "apply_pre_quantization", None)
        if fn is None:
            continue
        remaining, hooks = fn(
            model,
            state_dict,
            remaining,
            default_dtype=default_dtype,
            verboseLevel=verboseLevel,
        )
        if hooks:
            post_load.extend(hooks)
    return remaining, post_load

def _patch_marlin_fp8_bias():
    """
    Quanto's Marlin FP8 CUDA kernel currently ignores the bias argument.
    Add it back manually (in-place) so outputs stay correct on CUDA builds.
    """
    try:
        from optimum.quanto.tensor.weights.marlin.fp8 import qbits as marlin_fp8
    except Exception:
        return
    if getattr(marlin_fp8.MarlinF8QBytesLinearFunction, "_wan2gp_bias_patch", False):
        return

    orig_forward = marlin_fp8.MarlinF8QBytesLinearFunction.forward

    def forward_with_bias(ctx, input, other, bias=None):
        out = orig_forward(ctx, input, other, None)
        if bias is None:
            return out
        bias_to_add = bias
        if bias_to_add.device != out.device or bias_to_add.dtype != out.dtype:
            bias_to_add = bias_to_add.to(device=out.device, dtype=out.dtype)
        view_shape = [1] * out.ndim
        view_shape[-1] = bias_to_add.shape[0]
        bias_view = bias_to_add.view(*view_shape)
        out.add_(bias_view)
        return out

    marlin_fp8.MarlinF8QBytesLinearFunction.forward = staticmethod(forward_with_bias)  # type: ignore
    marlin_fp8.MarlinF8QBytesLinearFunction._wan2gp_bias_patch = True  # type: ignore
    marlin_fp8.MarlinF8QBytesLinearFunction._wan2gp_bias_orig = orig_forward  # type: ignore


_patch_marlin_fp8_bias()
