import subprocess
import os
import torch
from transformers import BitsAndBytesConfig, AutoConfig, AutoModelForCausalLM, LlavaNextForConditionalGeneration, LlavaForConditionalGeneration, PaliGemmaForConditionalGeneration, Idefics2ForConditionalGeneration
import spaces

os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"

# Install required package
def install_flash_attn():
    subprocess.run(
        "pip install flash-attn --no-build-isolation",
        env={"FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"},
        shell=True,
    )

# Architecture to model class mapping
ARCHITECTURE_MAP = {
    "LlavaNextForConditionalGeneration": LlavaNextForConditionalGeneration,
    "LlavaForConditionalGeneration": LlavaForConditionalGeneration,
    "PaliGemmaForConditionalGeneration": PaliGemmaForConditionalGeneration,
    "Idefics2ForConditionalGeneration": Idefics2ForConditionalGeneration,
    "AutoModelForCausalLM": AutoModelForCausalLM
}

# Function to get the model summary with caching and GPU support
@spaces.GPU
def get_model_summary(model_name):
    """
    Retrieve the model summary for the given model name.
    
    Args:
    model_name (str): The name of the model to retrieve the summary for.

    Returns:
    tuple: A tuple containing the model summary (str) and an error message (str), if any.
    """
    try:
        # Fetch the model configuration
        config = AutoConfig.from_pretrained(model_name)
        architecture = config.architectures[0]
        quantization_config = getattr(config, 'quantization_config', None)

        # Set up BitsAndBytesConfig if the model is quantized
        if quantization_config:
            bnb_config = BitsAndBytesConfig(
                load_in_4bit=quantization_config.get('load_in_4bit', False),
                load_in_8bit=quantization_config.get('load_in_8bit', False),
                bnb_4bit_compute_dtype=quantization_config.get('bnb_4bit_compute_dtype', torch.float16),
                bnb_4bit_quant_type=quantization_config.get('bnb_4bit_quant_type', 'nf4'),
                bnb_4bit_use_double_quant=quantization_config.get('bnb_4bit_use_double_quant', False),
                llm_int8_enable_fp32_cpu_offload=quantization_config.get('llm_int8_enable_fp32_cpu_offload', False),
                llm_int8_has_fp16_weight=quantization_config.get('llm_int8_has_fp16_weight', False),
                llm_int8_skip_modules=quantization_config.get('llm_int8_skip_modules', None),
                llm_int8_threshold=quantization_config.get('llm_int8_threshold', 6.0),
            )
        else:
            bnb_config = None

        # Get the appropriate model class from the architecture map
        model_class = ARCHITECTURE_MAP.get(architecture, AutoModelForCausalLM)
        
        # Load the model
        model = model_class.from_pretrained(
            model_name, config=bnb_config, trust_remote_code=True
        )

        # Move to device only if the model is not quantized
        if model and not quantization_config:
            model = model.to(torch.device("cuda" if torch.cuda.is_available() else "cpu"))

        model_summary = str(model) if model else "Model architecture not found."
        return model_summary, ""
    except ValueError as ve:
        return "", f"ValueError: {ve}"
    except EnvironmentError as ee:
        return "", f"EnvironmentError: {ee}"
    except Exception as e:
        return "", str(e)