import re
import os
import json
import torch

from safetensors.torch import load_file

from .specs import *
from .utils import compile_ignore, get_ignore, estimate_gpu_bytes
from .dequant_register import DEQUANT_KERNEL_TABLE
from .config_builder import create_quant_config_builder, resolve_quant_artifacts, to_dequant_spec



from safetensors import safe_open


class LazyTensorDict(dict):

    def __init__(self, base_tensors, tensor_index, input_dir):
        super().__init__(base_tensors)
        self.tensor_index = tensor_index
        self.input_dir = input_dir
        self.file_cache = {}

    def _load_tensor(self, name):

        file_name = self.tensor_index[name]
        path = os.path.join(self.input_dir, file_name)

        if file_name not in self.file_cache:
            self.file_cache[file_name] = safe_open(path, framework="pt")
        return self.file_cache[file_name].get_tensor(name)

    def __getitem__(self, key):
        try:
            return super().__getitem__(key)
        except KeyError:
            # 键不在字典中，检查是否在索引中
            if self.tensor_index and key in self.tensor_index:
                tensor = self._load_tensor(key)
                self[key] = tensor  # 现在可以安全地添加
                return tensor
            else:
                raise KeyError(key)
    
    def __contains__(self, key):

        if super().__contains__(key):
            return True

        if self.tensor_index and key in self.tensor_index:
            return True

        return False



def dequantize_(
    tensors: dict[str, torch.Tensor],
    qconfig: dict,
    custom_dequant_config,
    output_dtype=torch.bfloat16,
    device: torch.device = "cpu",
    **kwargs
):
    ignore_exact, ignore_regex = compile_ignore(
        get_ignore(qconfig),
        get_ignore(custom_dequant_config),
    )

    builder = create_quant_config_builder(
        qconfig,
        custom_dequant_config,
        model_type=kwargs.get("model_type")
    )

    outputs = {}
    gpu_cache = {}
    gpu_cache_size = 0
    max_gpu_cache_size = 16 * (1024**3)
    weight_names = list(tensors.keys())
    # for weight_name in tensors.keys():
    for weight_name in weight_names:
        rq = builder.resolve(weight_name)

        if rq is None:  # 包含scale的一定不会在这里
            assert "scale" not in weight_name, f"{weight_name}"
            outputs[weight_name] = tensors[weight_name]
            continue

        assert not rq.dynamic, "Dynamic quantization not supported in offline dequant"

        artifacts = resolve_quant_artifacts(weight_name, tensors, rq)

        if artifacts is None:
            # 跳过scale, shape等
            continue

        spec = to_dequant_spec(rq)
        if spec not in DEQUANT_KERNEL_TABLE:
            raise RuntimeError(f"No dequant kernel registered for {spec}")

        kernel = DEQUANT_KERNEL_TABLE[spec]

        gpu_cache_size += estimate_gpu_bytes(artifacts, rq)

        try:
            output = kernel(
                artifacts=artifacts,
                rq=rq,
                output_dtype=output_dtype,
                device=device,
            )
        except Exception as e:
            print(weight_name)
            raise e
        # 处理名称
        weight_name = weight_name.replace("weight_packed", "weight").replace("qweight", "weight")
        gpu_cache[weight_name] = output

        if gpu_cache_size >= max_gpu_cache_size:
            for k, v in gpu_cache.items():
                outputs[k] = v.cpu()
            gpu_cache.clear()
            gpu_cache_size = 0
    
    for k, v in gpu_cache.items():
        outputs[k] = v.cpu()
        
    gpu_cache.clear()
    gpu_cache_size = 0

    return outputs


def dequant(file_name: str,
            quant_config,
            custom_dequant_config,
            device: torch.device,
            model_type,
            tensor_index=None,
            input_dir=None):

    base_tensors = load_file(file_name)

    if tensor_index is not None:
        tensors = LazyTensorDict(base_tensors, tensor_index, input_dir)
    else:
        tensors = base_tensors

    return dequantize_(tensors, quant_config, custom_dequant_config, device=device, model_type=model_type)


if __name__ == "__main__":
    ...