"""diffusers_generate — ejecuta inferencia con un pipeline diffusers y retorna ImageGenResult."""

from __future__ import annotations

import sys
import os
import time
from typing import Any

sys.path.insert(0, os.path.dirname(__file__))

from generation_config import GenerationConfig
from image_gen_result import ImageGenResult
from genconfig_to_diffusers_kwargs import genconfig_to_diffusers_kwargs


def diffusers_generate(pipe: Any, cfg: GenerationConfig) -> ImageGenResult:
    """Ejecuta inferencia con un pipeline diffusers y retorna ImageGenResult.

    Convierte el GenerationConfig a kwargs via genconfig_to_diffusers_kwargs,
    crea un torch.Generator con la semilla configurada, mide duracion y pico
    de VRAM (si CUDA disponible). El campo meta del resultado incluye backend,
    modelo, sampler, seed y steps usados.

    Args:
        pipe:  Pipeline diffusers cargado (resultado de diffusers_load_pipeline).
               Debe ser callable: pipe(prompt=..., ...) -> objeto con .images[0].
        cfg:   Parametros de generacion. cfg.seed se usa para torch.Generator.
               cfg.model.name se incluye en meta. cfg.sampler se incluye en meta
               pero NO se aplica aqui — usar diffusers_set_scheduler antes si se
               quiere cambiar el sampler.

    Returns:
        ImageGenResult con image=PIL.Image, meta con keys backend/model/sampler/
        actual_steps/seed, duration_ms y vram_peak_mb (None si no hay CUDA).

    Raises:
        ImportError: Si torch o diffusers no estan instalados.
        RuntimeError: Si la inferencia falla (OOM, modelo incompatible, etc.).
    """
    try:
        import torch
    except ImportError as exc:
        raise ImportError(
            "diffusers_generate requiere torch. "
            "Instalar con: pip install torch"
        ) from exc

    # Determinar device del pipeline
    device = "cpu"
    if hasattr(pipe, "device"):
        device = str(pipe.device)

    # Medir VRAM solo en CUDA
    cuda_available = torch.cuda.is_available()
    if cuda_available:
        torch.cuda.reset_peak_memory_stats()

    # Construir kwargs desde GenerationConfig
    kwargs = genconfig_to_diffusers_kwargs(cfg)

    # Crear generator con semilla
    seed = cfg.seed if cfg.seed >= 0 else int(time.time()) % (2**32)
    generator = torch.Generator(device=device).manual_seed(seed)
    kwargs["generator"] = generator

    # Inferencia
    t0 = time.perf_counter()
    result = pipe(**kwargs)
    t1 = time.perf_counter()

    duration_ms = int((t1 - t0) * 1000)

    # VRAM peak
    vram_peak_mb: int | None = None
    if cuda_available:
        vram_peak_mb = torch.cuda.max_memory_allocated() // 1024 // 1024

    # Nombre del modelo
    model_name = cfg.model.name if cfg.model and hasattr(cfg.model, "name") else "unknown"

    meta: dict[str, Any] = {
        "backend": "diffusers",
        "model": model_name,
        "sampler": cfg.sampler,
        "actual_steps": cfg.steps,
        "seed": seed,
        "width": cfg.width,
        "height": cfg.height,
        "cfg_scale": cfg.cfg_scale,
    }

    return ImageGenResult(
        image=result.images[0],
        meta=meta,
        duration_ms=duration_ms,
        vram_peak_mb=vram_peak_mb,
    )
