# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

import json
import os

from compressed_tensors import __version__ as ct_version
from compressed_tensors.base import COMPRESSION_VERSION_NAME, QUANTIZATION_CONFIG_NAME
from compressed_tensors.entrypoints.convert import Converter
from compressed_tensors.utils.safetensors_load import find_config_path
from loguru import logger
from safetensors.torch import load_file, save_file


__all__ = [
    "validate_file",
    "convert_file",
    "write_checkpoint_quantization_config",
]


def write_checkpoint_quantization_config(
    save_directory: str | os.PathLike,
    converter: Converter,
):
    """
    Write the quantization config produced by `converter` into the model config
    file (config.json or params.json) in save_directory. This is called after
    the convert checkpoint pathway completes to record which quantization was
    applied. The quantization_config section is replaced entirely with the new
    config.

    :param save_directory: directory containing the model config file
    :param converter: Converter instance whose create_config() produces the
        updated quantization config
    """
    quant_config = converter.create_config()

    quant_config_data = quant_config.model_dump()
    quant_config_data[COMPRESSION_VERSION_NAME] = ct_version

    config_file_path = find_config_path(save_directory)
    if config_file_path is not None:
        with open(config_file_path, "r") as file:
            config_data = json.load(file)

        config_data[QUANTIZATION_CONFIG_NAME] = quant_config_data

        with open(config_file_path, "w") as file:
            json.dump(config_data, file, indent=2, sort_keys=True)

    else:
        logger.warning(
            f"Could not find config file in {save_directory}. Please add to config "
            f"{json.dumps(quant_config_data, indent=2, sort_keys=True)}"
        )


def validate_file(
    file_path: str | os.PathLike,
    converter: Converter,
):
    """
    Validate that each quantizable tensor in a safetensors file can be quantized.

    :param file_path: safetensors file to validate
    :param converter: converter we wish to apply to the checkpoint,
        e.g. conversion of some layers from some format to compressed-tensors
    """
    tensors = load_file(file_path)

    converter.validate(tensors)


def convert_file(
    file_path: str | os.PathLike,
    save_path: str | os.PathLike,
    converter: Converter,
) -> tuple[int, dict[str, str]]:
    """
    Convert tensors in a given safetensors file

    :param file_path: safetensors file to process
    :param save_path: save path of file with quantized weights
    :param converter: converter we wish to apply to the checkpoint,
        e.g. conversion of some layers from some format to compressed-tensors
    """
    tensors = load_file(file_path)

    converter.process(tensors)

    save_file(tensors, save_path)
    total_size = sum(tensor.nbytes for tensor in tensors.values())
    weight_map = {key: os.path.basename(save_path) for key in tensors.keys()}
    return total_size, weight_map
