import ctypes
import json
import math
import os
import struct
import warnings

import numpy as np
import torch
from safetensors import safe_open
from tqdm import tqdm

DTYPE_ELEM_SIZE = {
    "BF16": 2, "F16": 2, "F32": 4, "F64": 8,
    "I8": 1, "U8": 1, "BOOL": 1,
    "I16": 2, "U16": 2,
    "I32": 4, "U32": 4,
    "I64": 8, "U64": 8,
}

HADAMARD_CACHE = {}


# матрица Адамара (НЕ менять базис — см. предупреждение в шапке файла)
def get_h4():
    return torch.tensor([
        [1.0, 1.0, 1.0, -1.0],
        [1.0, 1.0, -1.0, 1.0],
        [1.0, -1.0, 1.0, 1.0],
        [-1.0, 1.0, 1.0, 1.0]
    ], dtype=torch.float32) * 0.5


def get_hadamard_matrix(group_size: int, device="cpu"):
    """
    строит матрицу Адамара размера group_size x group_size через
    рекурсивный kron() с H4. Работает КОРРЕКТНО только для group_size,
    являющегося степенью 4 (4, 16, 64, 256, 1024, ...)
    """
    dev_key = str(device)
    if (group_size, dev_key) in HADAMARD_CACHE:
        return HADAMARD_CACHE[(group_size, dev_key)]

    h = get_h4()
    cur_size = 4
    while cur_size < group_size:
        h = torch.kron(h, get_h4())
        cur_size *= 4

    if cur_size != group_size:
        raise ValueError(
            f"group_size={group_size} не является степенью 4 (4,16,64,256,...). "
            f"текущая реализация build-а Hadamard-матрицы через kron(H4,H4,...) "
            f"поддерживает только такие размеры."
            f"проверь значение 'convrot_groupsize' в .comfy_quant метаданных этого чекпоинта — "
            f"если оно реально другое, эту функцию нужно расширять под их схему, "
            f"а не подставлять сюда \"обычный\" Hadamard — это будет другой базис "
            f"и де-ротация окажется математически неверной."
        )

    h = h.to(device)
    HADAMARD_CACHE[(group_size, dev_key)] = h
    return h


def coerce_bool(value):
    """
    bool("false") == True в Python — ловушка для строковых JSON-флагов.
    """
    if isinstance(value, str):
        return value.strip().lower() not in ("false", "0", "", "no")
    return bool(value)


# low-level запись
def write_tensor_to_file(t: torch.Tensor, out_f):
    """
    zero-copy запись сырых байт тензора на диск.
    """
    t = t.contiguous().cpu()
    num_bytes = t.numel() * t.element_size()
    if num_bytes == 0:
        return
    buf = (ctypes.c_char * num_bytes).from_address(t.data_ptr())
    out_f.write(buf)


def decode_metadata_tensor(t: torch.Tensor) -> bytes:
    """
    numpy().view(np.uint8) переинтерпретирует те же самые байты без
    привязки к знаковости dtype — устойчиво к I8/U8/любому 1-байтовому
    представлению.
    """
    arr = t.cpu().numpy()
    if arr.dtype != np.uint8:
        arr = arr.view(np.uint8)
    return arr.tobytes()


def resolve_tensor_spec(key, shape, native_dtype_str, is_quantized, target_dtype_str, target_elem_size):
    """
    единая точка принятия решения "каким будет этот тензор на
    выходе" — используется и в проходе 1 (расчёт заголовка), и в
    проходе 2 (запись). Раньше эта логика была продублирована в двух
    местах и теоретически могла разойтись.
    """
    if is_quantized:
        return shape, target_dtype_str, target_elem_size
    if native_dtype_str in ("BF16", "F16", "F32"):
        return shape, target_dtype_str, target_elem_size
    elem_s = DTYPE_ELEM_SIZE.get(native_dtype_str, 4)
    return shape, native_dtype_str, elem_s


def stream_dequant_convrot(
        input_path: str,
        output_path: str,
        target_dtype: torch.dtype = torch.bfloat16,
        strict: bool = True,
):
    """
    strict=True (по умолчанию): при любой неоднозначности (нужна ли расротация,
    group_size не делит in_features и т.п.) - останавливается с ошибкой.
    strict=False: эвристика + warning.
    """
    if not os.path.exists(input_path):
        raise FileNotFoundError(f"Файл {input_path} не найден!")
    if os.path.abspath(input_path) == os.path.abspath(output_path):
        raise ValueError("input_path и output_path совпадают — это перезапишет исходник.")

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"используемое устройство вычислений: {device}")

    dtype_str_map = {
        torch.bfloat16: "BF16",
        torch.float16: "F16",
        torch.float32: "F32",
    }
    if target_dtype not in dtype_str_map:
        raise ValueError(f"неподдерживаемый target_dtype: {target_dtype}")
    target_dtype_str = dtype_str_map[target_dtype]
    target_elem_size = 2 if target_dtype in (torch.bfloat16, torch.float16) else 4

    print("--- сканирование архитектуры и расчёт заголовка ---")

    quant_configs = {}
    broken_configs = []
    tensor_specs = {}
    write_order = []

    with safe_open(input_path, framework="pt", device="cpu") as f:
        raw_meta = f.metadata()
        file_metadata = {str(k): str(v) for k, v in raw_meta.items()} if raw_meta is not None else {}

        keys = list(f.keys())
        all_keys_set = set(keys)

        for k in keys:
            if k.endswith(".comfy_quant"):
                base_name = k[:-len(".comfy_quant")]
                raw_bytes = decode_metadata_tensor(f.get_tensor(k))
                try:
                    quant_configs[base_name] = json.loads(raw_bytes.decode("utf-8"))
                except Exception as e:
                    # список проблемных слоёв фиксируется и печатаем предупреждение,
                    # т.к. для таких слоёв is_convrot придётся определять по эвристике
                    broken_configs.append(base_name)
                    warnings.warn(f"не удалось распарсить .comfy_quant для '{base_name}': {e}")

        zero_point_hits = [
            k for k in keys
            if k.endswith(".weight_zero_point")
        ]
        if zero_point_hits:
            raise NotImplementedError(
                f"найдены тензоры zero-point (асимметричная квантизация), "
                f"например: {zero_point_hits[:10]}."
                f"формула деквантизации в этом скрипте (value = int8 * scale) "
                f"рассчитана на симметричную квантизацию без zero-point. "
                f"раскодировать такие слои по текущей формуле — получить смещённые веса. "
                f"нужно явно реализовать 'value = (int8 - zero_point) * scale' перед тем, как продолжать."
            )

        for k in keys:
            if k.endswith(".comfy_quant") or k.endswith(".weight_scale"):
                continue

            base_name = k[:-len(".weight")] if k.endswith(".weight") else None
            shape = list(f.get_slice(k).get_shape())
            t_slice = f.get_slice(k)
            native_dt = t_slice.get_dtype()

            is_quantized = bool(
                base_name and (base_name in quant_configs or f"{base_name}.weight_scale" in all_keys_set)
            )
            tensor_specs[k] = resolve_tensor_spec(
                k, shape, native_dt, is_quantized, target_dtype_str, target_elem_size
            )
            write_order.append(k)

    n_quant = sum(
        1 for k in write_order
        if k.endswith(".weight") and k[:-len(".weight")] in quant_configs
    )
    print(f"всего слоёв: {len(write_order)}")
    print(f"обнаружено квантованных слоёв (с явным конфигом): {n_quant}")
    if broken_configs:
        print(f"ВНИМАНИЕ: {len(broken_configs)} слоёв с нечитаемым .comfy_quant: {broken_configs[:5]}...")

    header = {}
    current_offset = 0
    for k in write_order:
        shape, dt_str, elem_size = tensor_specs[k]
        num_elements = math.prod(shape) if shape else 1
        num_bytes = num_elements * elem_size
        header[k] = {
            "dtype": dt_str,
            "shape": shape,
            "data_offsets": [current_offset, current_offset + num_bytes]
        }
        current_offset += num_bytes

    file_metadata.setdefault("format", "pt")
    header["__metadata__"] = file_metadata
    header_json = json.dumps(header, separators=(",", ":")).encode("utf-8")

    pad_len = (8 - ((8 + len(header_json)) % 8)) % 8
    header_json += b" " * pad_len
    header_len = len(header_json)

    print(f"размер распакованной модели составит: {current_offset / (1024 ** 3):.2f} GB")
    print("--- потоковая конвертация и запись на диск ---")

    tmp_output_path = output_path + ".tmp"
    stats = {"rotated": 0, "skipped_no_divide": 0, "heuristic_fallback": 0}

    try:
        with open(tmp_output_path, "wb", buffering=16 * 1024 * 1024) as out_f:
            out_f.write(struct.pack("<Q", header_len))
            out_f.write(header_json)

            with safe_open(input_path, framework="pt", device="cpu") as f:
                for k in tqdm(write_order, desc="восстановление весов"):
                    base_name = k[:-len(".weight")] if k.endswith(".weight") else None
                    scale_key = f"{base_name}.weight_scale" if base_name else None

                    if base_name and (base_name in quant_configs or scale_key in all_keys_set):
                        w_int8 = f.get_tensor(k)
                        scale = f.get_tensor(scale_key)
                        cfg = quant_configs.get(base_name, {})

                        w_gpu = w_int8.to(device=device, dtype=torch.float32)
                        s_gpu = scale.to(device=device, dtype=torch.float32)

                        if s_gpu.numel() != w_gpu.shape[0]:
                            raise ValueError(
                                f"'{k}': scale имеет {s_gpu.numel()} элементов, "
                                f"а out_features={w_gpu.shape[0]}. "
                                f"Предположение о поканальном (per-out-channel) масштабе не выполняется — "
                                f"формула дальше даст неверный broadcast."
                            )
                        s_gpu = s_gpu.view(s_gpu.shape[0], *([1] * (w_gpu.ndim - 1)))
                        w_float = w_gpu * s_gpu

                        # по документации Comfy-Org/comfy-quants формат маркера
                        # это {"convrot": true/false, "convrot_groupsize": N, "per_row": true}
                        # (или, в int8_tensorwise, ключи convrot/convrot_groupsize просто
                        # опускается, если слой не был повёрнут).
                        # значит, если маркер для base_name успешно распарсился,
                        # то отсутствие ключа "convrot" внутри него — не двусмысленность,
                        # а штатное "convrot": false (writer сам его так кодирует, когда
                        # in_features % group_size != 0).
                        # Поэтому здесь угадывать нечего — просто .get(..., False).
                        #
                        # неопределённость — только когда base_name ВООБЩЕ
                        # отсутствует в quant_configs (маркер не сохранился или не
                        # распарсился), хотя scale при этом есть: тогда мы не можем
                        # узнать, был ли слой повёрнут.
                        if base_name in quant_configs:
                            is_convrot = coerce_bool(quant_configs[base_name].get("convrot", False))
                        else:
                            msg = (
                                f"'{k}': есть '{scale_key}', но нет распознанного "
                                f".comfy_quant маркера для '{base_name}' (файл сохранён без "
                                f"--save-quant-metadata, либо JSON не распарсился — "
                                f"см. предупреждения выше). неизвестно, была ли применена "
                                f"ConvRot-расротация к этому слою."
                            )
                            if strict:
                                raise ValueError(
                                    msg + " запусти с strict=False, чтобы принять "
                                          "заведомо небезопасное допущение is_convrot=False."
                                )
                            warnings.warn(msg + " используется допущение: is_convrot=False.")
                            is_convrot = False
                            stats["heuristic_fallback"] += 1

                        if is_convrot and w_float.ndim == 2:
                            group_size = cfg.get("convrot_groupsize", 256)
                            out_features, in_features = w_float.shape

                            if in_features % group_size == 0:
                                H = get_hadamard_matrix(group_size, device=device)
                                w_grouped = w_float.view(out_features, -1, group_size)
                                w_unrot = torch.matmul(w_grouped, H)
                                w_float = w_unrot.reshape(out_features, in_features)
                                stats["rotated"] += 1
                            else:
                                msg = (
                                    f"'{k}': in_features={in_features} не делится на "
                                    f"group_size={group_size} — расротация ПРОПУЩЕНА, "
                                    f"веса останутся в повёрнутом виде."
                                )
                                if strict:
                                    raise ValueError(msg + " запусти с strict=False, чтобы всё же продолжить.")
                                warnings.warn(msg)
                                stats["skipped_no_divide"] += 1

                        tensor_to_save = w_float.to(dtype=target_dtype, device="cpu").contiguous()
                        del w_int8, scale, w_gpu, s_gpu, w_float
                    else:
                        t = f.get_tensor(k)
                        if t.is_floating_point():
                            tensor_to_save = t.to(dtype=target_dtype, device="cpu").contiguous()
                        else:
                            tensor_to_save = t.contiguous().cpu()
                        del t

                    write_tensor_to_file(tensor_to_save, out_f)
                    del tensor_to_save

            if torch.cuda.is_available():
                torch.cuda.empty_cache()

        os.replace(tmp_output_path, output_path)
    except Exception:
        if os.path.exists(tmp_output_path):
            os.remove(tmp_output_path)
        raise

    print("\n--- сводка ---")
    print(f"де-ротировано слоёв: {stats['rotated']}")
    print(f"пропущено разворотов (group_size не делит in_features): {stats['skipped_no_divide']}")
    print(f"слоёв, где is_convrot определён эвристикой (не из конфига): {stats['heuristic_fallback']}")
    if stats["skipped_no_divide"] or stats["heuristic_fallback"]:
        print("ВНИМАНИЕ: см. предупреждения выше — часть слоёв обработана неоднозначно.")

    # файл должен как минимум открываться и
    # содержать ровно то же число тензоров
    with safe_open(output_path, framework="pt", device="cpu") as f_check:
        n_out = len(f_check.keys())
        if n_out != len(write_order):
            raise RuntimeError(
                f"проверка после записи не прошла: в выходном файле {n_out} тензоров, "
                f"ожидалось {len(write_order)}."
            )

    print(f"\n[успешно] модель полностью восстановлена и сохранена в: {output_path}")


if __name__ == "__main__":
    INPUT_FILE = "/content/SinoX_Vision_Krea2.safetensors"
    OUTPUT_FILE = "/content/SinoX_Vision_Krea2_bf16.safetensors"

    stream_dequant_convrot(
        input_path=INPUT_FILE,
        output_path=OUTPUT_FILE,
        target_dtype=torch.bfloat16,
        strict=True,  # поставь False, если готов принять эвристику
    )