from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
)
from optimum.onnxruntime import ORTModelForCausalLM

from torch import bfloat16 as torch_bfloat16
from torch import float16 as torch_float16
from torch import float32 as torch_float32
from torch import compile as torch_compile

from comfy.model_management import (
    get_torch_device,
    should_use_fp16,
    should_use_bf16,
    is_device_cuda,
)

from .utility import (
    ModelType,
    QuantizationType,
    QuantizationPackage,
    get_quantization_package,
    is_base_model,
    check_torch_version_is_enough,
)

from peft import PeftModel


def get_torch_dtype():
    dev = get_torch_device()

    if should_use_bf16(device=dev):
        req_torch_dtype = torch_bfloat16
    elif should_use_fp16(device=dev):
        req_torch_dtype = torch_float16
    else:
        req_torch_dtype = torch_float32

    return req_torch_dtype


def get_bitsandbytes_config(type: QuantizationType):
    from transformers import BitsAndBytesConfig

    bnb_config = None

    if type == QuantizationType.EightBit:
        bnb_config = BitsAndBytesConfig(load_in_8bit=True)
    elif type == QuantizationType.FourBit:
        bnb_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_use_double_quant=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype=get_torch_dtype(),
        )

    return bnb_config


def get_model_from_base(
    model_name: str, required_torch_dtype, type: QuantizationType, is_acceleration: bool
):
    model_configs = {}

    model_configs["device_map"] = "auto"
    model_configs["torch_dtype"] = required_torch_dtype

    quant_pack = get_quantization_package()

    if quant_pack == QuantizationPackage.BITSANDBYTES:
        quant_conf = get_bitsandbytes_config(type)

        model_configs["quantization_config"] = quant_conf

    if is_acceleration and check_torch_version_is_enough(2, 1):
        model_configs["attn_implementation"] = "sdpa"

    model = AutoModelForCausalLM.from_pretrained(model_name, **model_configs)

    if quant_pack == QuantizationPackage.QUANTO:
        from optimum.quanto import qfloat8, qint8, qint4, quantize, freeze

        if type == QuantizationType.EightBit:
            quantize(model, weights=qint8, activations=qfloat8, exclude="lm_head")
        elif type == QuantizationType.EightFloat:
            quantize(model, weights=qfloat8, activations=qfloat8, exclude="lm_head")
        elif type == QuantizationType.FourBit:
            quantize(model, weights=qint4, activations=qfloat8, exclude="lm_head")

        freeze(model)

    return model


def get_model_from_lora(
    model_name: str, required_torch_dtype, type: QuantizationType, is_acceleration: bool
):
    model = get_model_from_base(model_name, required_torch_dtype, type, is_acceleration)

    model.config.forced_decoder_ids = None
    model.config.suppress_tokens = []

    model = PeftModel.from_pretrained(model, model_name, is_trainable=False)

    return model


def get_model(model_name: str, type: QuantizationType, is_acceleration: bool):
    req_torch_dtype = get_torch_dtype()

    if is_base_model(model_name):
        model = get_model_from_base(model_name, req_torch_dtype, type, is_acceleration)
    else:
        model = get_model_from_lora(model_name, req_torch_dtype, type, is_acceleration)

    torch_compile(model)

    return model


def get_tokenizer(model_name: str):
    # use the fast implementation of the tokenizer if possible
    try:
        tokenizer = AutoTokenizer.from_pretrained(model_name)
    except:
        tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)

    return tokenizer


def get_model_tokenizer(model_path: str, type: ModelType, quant_type: QuantizationType):
    """
    transformers -> ONNX operation brokes often
    """
    if type == ModelType.ONNX:
        model_configs = {}

        if is_device_cuda(get_torch_device()):
            model_configs["provider"] = "CUDAExecutionProvider"

        model = ORTModelForCausalLM.from_pretrained(model_path, **model_configs)
    elif type == ModelType.BETTERTRANSFORMER:
        model = get_model(model_path, quant_type, is_acceleration=True)
    elif type == ModelType.DEFAULT:
        model = get_model(model_path, quant_type, is_acceleration=False)

    tokenizer = get_tokenizer(model_path)

    return (model, tokenizer)


__all__ = [get_model_tokenizer]