import json
from pathlib import Path

CONFIG_FILE = Path(".QuantTensor_config.json")
__version__ = "0.0.1"

# === Chargement ou création de la config ===
if CONFIG_FILE.exists():
    with open(CONFIG_FILE, "r") as f:
        config = json.load(f)
else:
    config = None

# Vérification si on doit relancer les tests
if config is None or config.get("version") != __version__:
    print("[QuantTensor] Première exécution ou mise à jour détectée : test des capacités Triton...")

    # Test réel de Triton (import conditionnel)
    try:
        import triton
        import triton.language as tl
        import torch
        print(f"[QuantTensor] Triton est fonctionnel.")

        try:
            pass # Tester le premier kernel de Quant
            
            print(f" [X] Kernels de quantifications Q4_1")
            Q4_1__triton_quant = True
        except Exception as e :
            
            print(f" [ ] Kernels de quantifications Q4_1 :\n{e}")
            Q4_1__triton_quant = False
        
        try:
            pass # Tester le premier kernel de DeQuant

            print(f" [X] Kernels de déquantifications Q4_1")
            Q4_1__triton_dequant = True
        except Exception as e :

            print(f" [ ] Kernels de déquantifications Q4_1 :\n{e}")
            Q4_1__triton_dequant = False
        
        try:
            pass # Tester kernel de GEMM même si les précédents passent pas

            print(f" [X] Kernels de fused-GEMM Q4_1")
            Q4_1__triton_fused_gemm = True
        except Exception as e :

            print(f" [ ] Kernels de fused-GEMM Q4_1 :\n{e}")
            Q4_1__triton_fused_gemm = False
        
        # Et Etc....

    except Exception as e :
        print(f"[QuantTensor] La bibliothèque Triton ne fonctionne pas :\n{e}")
        Q4_1__triton_dequant = False
        Q4_1__triton_quant = False
        Q4_1__triton_fused_gemm = False
        

    config = {
        "version": __version__,
        "Q4_1__triton_dequant": Q4_1__triton_dequant,
        "Q4_1__triton_quant": Q4_1__triton_quant,
        "Q4_1__fused_gemm": Q4_1__triton_fused_gemm
    }

    with open(CONFIG_FILE, "w") as f:
        json.dump(config, f, indent=2)



# === Binding des fonctions au moment de l'import ===
if config["triton_dequant"]:
    from ops import dequantize_q4_1_triton as dequantize_q4_1
else:
    from ops import dequantize_q4_1_pytorch as dequantize_q4_1

if config["triton_quant"]:
    from ops import quantize_q4_1_triton as dequantize_q4_1
else:
    from ops import quantize_q4_1_pytorch as dequantize_q4_1

if config["q4_1_fused_gemm"]:
    from ops import q4_1_fused_gemm_triton as q4_1_fused_gemm
else:
    from ops import q4_1_fused_gemm_pytorch as q4_1_fused_gemm

# =================

class QuantTensor:
    def __init__(self, format: str, qs: torch.Tensor, d: torch.Tensor, 
                 m: torch.Tensor, original_shape: torch.Size):
        self.format = format
        self.qs = qs
        self.d = d
        self.m = m
        self.original_shape = original_shape

        self.dtype = d.dtype
        self.device = d.device

    def numel(self) -> int:
        return self.original_shape.numel()

    @property
    def nbytes(self) -> int:
        total = 0
        for attr in self.__dict__.values():
            if isinstance(attr, torch.Tensor):
                total += attr.numel() * attr.element_size()
        return total

    def dequantize(self):
        if self.format == "Q4_1":
            w = dequantize_q4_1(self.qs, self.d, self.m)
            return w.view(self.original_shape)
        else:
            raise NotImplementedError(f"Format {self.format} non supporté")

    def to(self, device=None, dtype=None):
        if device is not None:
            self.device = device
            self.qs = self.qs.to(device=device)
            self.d = self.d.to(device=device)
            self.m = self.m.to(device=device)

        if dtype is not None:
            self.dtype = dtype

        return self