from dataclasses import dataclass, field
import sys


@dataclass
class ExllamaConfig:
    max_seq_len: int
    gpu_split: str = None
    cache_8bit: bool = False


class ExllamaModel:
    def __init__(self, exllama_model, exllama_cache):
        self.model = exllama_model
        self.cache = exllama_cache
        self.config = self.model.config


def load_exllama_model(model_path, exllama_config: ExllamaConfig):
    try:
        from exllamav2 import (
            ExLlamaV2Config,
            ExLlamaV2Tokenizer,
            ExLlamaV2,
            ExLlamaV2Cache,
            ExLlamaV2Cache_8bit,
        )
    except ImportError as e:
        print(f"Error: Failed to load Exllamav2. {e}")
        sys.exit(-1)

    exllamav2_config = ExLlamaV2Config()
    exllamav2_config.model_dir = model_path
    exllamav2_config.prepare()
    exllamav2_config.max_seq_len = exllama_config.max_seq_len
    exllamav2_config.cache_8bit = exllama_config.cache_8bit

    exllama_model = ExLlamaV2(exllamav2_config)
    tokenizer = ExLlamaV2Tokenizer(exllamav2_config)

    split = None
    if exllama_config.gpu_split:
        split = [float(alloc) for alloc in exllama_config.gpu_split.split(",")]
    exllama_model.load(split)

    cache_class = ExLlamaV2Cache_8bit if exllamav2_config.cache_8bit else ExLlamaV2Cache
    exllama_cache = cache_class(exllama_model)
    model = ExllamaModel(exllama_model=exllama_model, exllama_cache=exllama_cache)

    return model, tokenizer