import gc
import logging
from pathlib import Path

import torch
from huggingface_hub import snapshot_download
from transformers import AutoProcessor, BitsAndBytesConfig, Qwen3VLForConditionalGeneration

import folder_paths

from .model_registry import dtype_from_name


logger = logging.getLogger("FL_VLM")

_MODEL_CACHE = {}


def get_vlm_models_dir():
    model_dir = Path(folder_paths.models_dir) / "VLM"
    model_dir.mkdir(parents=True, exist_ok=True)
    return model_dir


def register_vlm_folder():
    model_dir = str(get_vlm_models_dir())
    supported = {".safetensors", ".bin", ".json", ".txt", ".model"}
    if "VLM" not in folder_paths.folder_names_and_paths:
        folder_paths.folder_names_and_paths["VLM"] = ([model_dir], supported)
        return

    paths, extensions = folder_paths.folder_names_and_paths["VLM"]
    if model_dir not in paths:
        paths.append(model_dir)
    extensions.update(supported)


def _safe_model_dir_name(model_id):
    return model_id.replace("/", "--")


def _quantization_config(quantization):
    if quantization == "4bit":
        return BitsAndBytesConfig(load_in_4bit=True)
    if quantization == "8bit":
        return BitsAndBytesConfig(load_in_8bit=True)
    return None


def _device_map(device_map):
    if device_map == "cuda":
        return {"": "cuda:0"}
    if device_map == "cpu":
        return {"": "cpu"}
    return device_map


def clear_model_cache():
    _MODEL_CACHE.clear()
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()


def load_vlm(model_id, dtype, quantization, attention, device_map, force_reload=False):
    register_vlm_folder()
    cache_key = (model_id, dtype, quantization, attention, device_map)

    if force_reload:
        _MODEL_CACHE.pop(cache_key, None)
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

    if cache_key in _MODEL_CACHE:
        logger.info("Using cached VLM model: %s", model_id)
        return _MODEL_CACHE[cache_key]

    model_path = get_vlm_models_dir() / _safe_model_dir_name(model_id)
    if not model_path.exists() or not any(model_path.iterdir()):
        logger.info("Downloading VLM model %s to %s", model_id, model_path)
        snapshot_download(repo_id=model_id, local_dir=str(model_path), local_dir_use_symlinks=False)

    torch_dtype = dtype_from_name(dtype)
    quant_config = _quantization_config(quantization)

    load_kwargs = {
        "torch_dtype": torch_dtype,
        "device_map": _device_map(device_map),
        "attn_implementation": attention,
    }
    if quant_config is not None:
        load_kwargs["quantization_config"] = quant_config

    logger.info("Loading VLM model from %s", model_path)
    model = Qwen3VLForConditionalGeneration.from_pretrained(str(model_path), **load_kwargs)
    processor = AutoProcessor.from_pretrained(str(model_path))

    model_bundle = {
        "backend": "qwen3_vl",
        "model": model,
        "processor": processor,
        "model_id": model_id,
        "model_path": str(model_path),
        "dtype": dtype,
        "quantization": quantization,
        "attention": attention,
        "device_map": device_map,
    }
    _MODEL_CACHE[cache_key] = model_bundle
    return model_bundle