"""Load Image21-INT4 or the BF16 source on CUDA, Apple Silicon, or CPU.

SDNQ stores UINT4 weights and dequantizes them with ordinary PyTorch operators
unless a CUDA/XPU matmul kernel is requested. This loader does not request that
kernel, so CUDA and MPS execute the same eager implementation. Bitsandbytes is
not used.
"""
import json
import os
from pathlib import Path

os.environ.setdefault('DIFFUSERS_SDNQ_TRANSFORMERS', '1')

import torch

from scripts.device import choose_runtime


def configure_caches():
    root = Path(__file__).resolve().parents[1]
    os.environ.setdefault('HF_HOME', str(root / '.cache' / 'huggingface'))
    os.environ.setdefault('NUMBA_CACHE_DIR', str(root / '.cache' / 'numba'))


def detect_device(requested=None):
    if requested:
        device = torch.device(requested)
        if device.type == 'cuda' and not torch.cuda.is_available():
            raise RuntimeError('CUDA was requested but is not available')
        if device.type == 'mps' and not torch.backends.mps.is_available():
            raise RuntimeError('MPS was requested but is not available')
        return device
    if torch.cuda.is_available():
        return torch.device('cuda')
    if torch.backends.mps.is_available():
        return torch.device('mps')
    return torch.device('cpu')


def accelerator_bytes(device):
    if device.type == 'cuda':
        return torch.cuda.get_device_properties(device).total_memory
    import psutil
    return psutil.virtual_memory().total


def read_quantization(model):
    path = Path(model) / 'transformer' / 'config.json'
    if not path.is_file():
        return {}
    return json.loads(path.read_text(encoding='utf-8')).get('quantization_config') or {}


def quantization_dict(module):
    config = getattr(module, 'config', None)
    raw = getattr(config, 'quantization_config', None)
    if raw is None:
        return {}
    if hasattr(raw, 'to_dict'):
        return raw.to_dict()
    if isinstance(raw, dict):
        return raw
    return dict(raw)


def _config_value(value):
    value = getattr(value, 'value', value)
    return str(value).split('.')[-1].lower()


def _require_int4(module, name):
    quant = quantization_dict(module)
    method = _config_value(quant.get('quant_method', ''))
    dtype = _config_value(quant.get('weights_dtype', ''))
    if method != 'sdnq' or dtype != 'uint4':
        raise ValueError(f'{name} is not an SDNQ UINT4 checkpoint: {method} {dtype}')
    if quant.get('use_quantized_matmul'):
        raise ValueError(f'{name} enables quantized matmul, which is not part of the portable INT4 runtime')


def _dtype(name):
    override = os.environ.get('IMAGE21_DTYPE')
    return getattr(torch, override or name)


def _enable_group(module, device, offload_type):
    from diffusers.hooks.group_offloading import apply_group_offloading
    kwargs = dict(onload_device=device, offload_device=torch.device('cpu'),
                  offload_type=offload_type, use_stream=False, non_blocking=False)
    if offload_type == 'block_level':
        kwargs['num_blocks_per_group'] = 1
    if hasattr(module, 'enable_group_offload'):
        module.enable_group_offload(**kwargs)
    else:
        apply_group_offloading(module, **kwargs)


def apply_offload(pipe, spec, device):
    mode = spec['offload']
    if mode == 'resident':
        pipe.to(device)
    elif mode == 'model':
        if device.type != 'cuda':
            raise ValueError('Model CPU offload is the CUDA path; Apple Silicon uses resident MPS')
        pipe.enable_model_cpu_offload(gpu_id=device.index or 0)
    elif mode == 'group':
        if device.type != 'cuda':
            raise ValueError('Group offload in this loader targets CUDA; Apple Silicon uses resident MPS')
        # The transformer is the activation-heavy component, so it moves one block at a time.
        # The text encoder moves one linear or embedding at a time: a whole decoder block's
        # dequantized MLP does not leave enough room under an 8GB cap.
        _enable_group(pipe.transformer, device, 'block_level')
        _enable_group(pipe.text_encoder, device, 'leaf_level')
        # A 1024 decode keeps several GiB live beside the denoiser cache. On an 8GB
        # board that peak does not fit, and tiling changes pixels. Decode on CPU.
        pipe.vae.to('cpu')
        decode = pipe.vae.decode
        encode_image = pipe._encode_vae_image

        def _decode_on_cpu(latents, *args, **kwargs):
            if torch.is_tensor(latents):
                latents = latents.detach().to('cpu')
            return decode(latents, *args, **kwargs)

        def _encode_on_cpu(image, generator):
            device = image.device
            encoded = encode_image(image.detach().to('cpu'), generator)
            return encoded.to(device=device, dtype=image.dtype)

        pipe.vae.decode = _decode_on_cpu
        pipe._encode_vae_image = _encode_on_cpu
    else:
        raise ValueError(f'Unknown offload mode: {mode}')
    pipe.image21_runtime = dict(spec, device=str(device))
    return pipe


def load_pipeline(model, offload='auto', device=None, local_files_only=False):
    """Load a local directory or a Hub repository id."""
    configure_caches()
    quant = read_quantization(model) if Path(model).exists() else {'quant_method': 'sdnq', 'weights_dtype': 'uint4'}
    method = str(quant.get('quant_method', '')).lower()
    if method == 'sdnq':
        import sdnq  # noqa: F401  (registers the quantizer with diffusers and transformers)
    elif quant:
        raise ValueError(f'Unsupported quantization config: {quant}')
    from diffusers import QwenImage21Pipeline
    device = detect_device(device)
    spec = choose_runtime(device.type, accelerator_bytes(device))
    if offload != 'auto':
        if offload not in ('model', 'group', 'resident'):
            raise ValueError('offload must be auto, model, group, or resident')
        spec = dict(spec, offload=offload)
    dtype = _dtype(spec['dtype'])
    pipe = QwenImage21Pipeline.from_pretrained(str(model), dtype=dtype, local_files_only=local_files_only)
    if method == 'sdnq':
        _require_int4(pipe.transformer, 'transformer')
        _require_int4(pipe.text_encoder, 'text_encoder')
        if quantization_dict(pipe.vae):
            raise ValueError('VAE is expected to stay in floating point')
    apply_offload(pipe, spec, device)
    return pipe