from __future__ import annotations

import gc
import time
import types
from typing import Any, Dict, Optional

import torch


def vram_snapshot() -> Dict[str, float]:
    if not torch.cuda.is_available():
        return {}
    snap = {}
    allocs = []
    peaks = []
    for i in range(torch.cuda.device_count()):
        alloc = torch.cuda.memory_allocated(i) / 1e9
        reserved = torch.cuda.memory_reserved(i) / 1e9
        peak = torch.cuda.max_memory_allocated(i) / 1e9
        snap[f'gpu{i}_alloc_gb'] = float(alloc)
        snap[f'gpu{i}_reserved_gb'] = float(reserved)
        snap[f'gpu{i}_peak_gb'] = float(peak)
        allocs.append(alloc)
        peaks.append(peak)
    snap['sum_alloc_gb'] = float(sum(allocs))
    snap['max_alloc_gb'] = float(max(allocs)) if allocs else 0.0
    snap['max_peak_gb'] = float(max(peaks)) if peaks else 0.0
    return snap


def clear_cuda():
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()


def infer_input_device(model) -> torch.device:
    try:
        emb = model.get_input_embeddings()
        if emb is not None:
            return next(emb.parameters()).device
    except Exception:
        pass
    try:
        for p in model.parameters():
            if not getattr(p, 'is_meta', False):
                return p.device
    except Exception:
        pass
    return torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')


def attach_diffusiongemma_block_diffusion(
    bot: Any,
    *,
    model_id: str = 'google/diffusiongemma-26B-A4B-it',
    device_map: str = 'auto',
    dtype: str = 'auto',
    trust_remote_code: bool = False,
    default_max_new_tokens: int = 512,
    verbose: bool = True,
):
    try:
        from transformers import AutoProcessor, DiffusionGemmaForBlockDiffusion
        model_cls = DiffusionGemmaForBlockDiffusion
        model_class_name = 'DiffusionGemmaForBlockDiffusion'
    except Exception:
        from transformers import AutoProcessor, AutoModelForMultimodalLM
        model_cls = AutoModelForMultimodalLM
        model_class_name = 'AutoModelForMultimodalLM'

    processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=trust_remote_code)
    kwargs = {'device_map': device_map, 'trust_remote_code': trust_remote_code}
    if dtype is not None:
        kwargs['dtype'] = dtype
    model = model_cls.from_pretrained(model_id, **kwargs)
    model.eval()
    input_device = infer_input_device(model)

    bot.model = model
    bot.processor = processor
    bot.tokenizer = getattr(processor, 'tokenizer', None)
    bot.input_device = input_device
    bot.model_loaded = True
    bot.model_id = model_id

    def diffusiongemma_generate_answer(self, system_prompt: str, user_prompt: str, max_new_tokens: Optional[int] = None, **generation_kwargs):
        t0 = time.perf_counter()
        model = self.model
        processor = self.processor
        device = getattr(self, 'input_device', None) or infer_input_device(model)
        messages = []
        if system_prompt:
            messages.append({'role': 'system', 'content': str(system_prompt)})
        messages.append({'role': 'user', 'content': str(user_prompt)})
        inputs = processor.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_dict=True, return_tensors='pt')
        if hasattr(inputs, 'to'):
            inputs = inputs.to(device)
        else:
            inputs = {k: (v.to(device) if torch.is_tensor(v) else v) for k, v in inputs.items()}
        input_len = int(inputs['input_ids'].shape[-1]) if 'input_ids' in inputs else 0
        gen_kwargs = dict(inputs)
        gen_kwargs['max_new_tokens'] = int(max_new_tokens or default_max_new_tokens)
        gen_kwargs.update(generation_kwargs)
        clear_cuda()
        with torch.inference_mode():
            outputs = model.generate(**gen_kwargs)
        out_ids = outputs[0] if isinstance(outputs, torch.Tensor) else outputs.sequences[0]
        try:
            text = processor.decode(out_ids[input_len:], skip_special_tokens=True).strip()
        except Exception:
            text = processor.decode(out_ids, skip_special_tokens=True).strip()
        return {'ran': True, 'answer': text, 'answer_raw': text, 'input_tokens': input_len, 'latency_s': float(time.perf_counter() - t0), 'model_id': model_id, 'adapter': 'NZFC DiffusionGemma adapter', 'vram': vram_snapshot()}

    bot.generate_answer = types.MethodType(diffusiongemma_generate_answer, bot)
    bot.nzfc_diffusiongemma_profile = {
        'version': 'v1.2.5a-runtime-assets-final',
        'base_model': model_id,
        'adapter': 'DiffusionGemma adapter',
        'claim_boundary': 'Marketing title uses Infinite-Context; technical layer is external evidence context, not native unlimited context.'
    }
    if verbose:
        print('[NZFC DiffusionGemma][OK] attached')
        print(bot.nzfc_diffusiongemma_profile)
    return bot.nzfc_diffusiongemma_profile