"""Model loading and caching for PixelGen."""

import logging
import os
from typing import Optional

import torch
from huggingface_hub import hf_hub_download

from .model_info import MODEL_CONFIGS, TEXT_ENCODER_CONFIGS, AVAILABLE_MODELS

logger = logging.getLogger("FL_PixelGen")

LOADED_MODELS_CACHE = {}
LOADED_ENCODERS_CACHE = {}

DTYPE_MAP = {
    "bfloat16": torch.bfloat16,
    "float16": torch.float16,
    "float32": torch.float32,
}


class PixelGenLoader:
    @staticmethod
    def load_model(model_name: str, device: str = "cuda", dtype: str = "bfloat16",
                   force_reload: bool = False):
        cache_key = f"{model_name}_{device}_{dtype}"
        if not force_reload and cache_key in LOADED_MODELS_CACHE:
            logger.info(f"Using cached model: {model_name}")
            return LOADED_MODELS_CACHE[cache_key]

        from src.models.transformer.JiT_T2I import JiT_T2I

        torch_dtype = DTYPE_MAP.get(dtype, torch.bfloat16)

        # Resolve checkpoint path
        if model_name in MODEL_CONFIGS:
            config = MODEL_CONFIGS[model_name]
            ckpt_path = PixelGenLoader._resolve_checkpoint(
                config["repo_id"], config["filename"]
            )
            denoiser_config = config["denoiser_config"]
        elif model_name in AVAILABLE_MODELS:
            ckpt_path = AVAILABLE_MODELS[model_name]
            denoiser_config = MODEL_CONFIGS.get(
                "PixNerd-XXL-P16-T2I", {}
            ).get("denoiser_config", {})
        else:
            raise ValueError(f"Unknown model: {model_name}")

        logger.info(f"Loading denoiser from: {ckpt_path}")

        # Instantiate denoiser
        denoiser = JiT_T2I(**denoiser_config)

        # Load EMA weights
        ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
        state_dict = ckpt.get("state_dict", ckpt)
        prefix = "ema_denoiser."
        loaded = 0
        for k, v in denoiser.state_dict().items():
            key = prefix + k
            if key in state_dict:
                v.copy_(state_dict[key])
                loaded += 1
            else:
                logger.warning(f"Missing key in checkpoint: {key}")
        logger.info(f"Loaded {loaded} weight tensors from checkpoint")

        denoiser = denoiser.to(device=device, dtype=torch_dtype)
        denoiser.eval()

        result = {
            "denoiser": denoiser,
            "device": device,
            "dtype": torch_dtype,
            "model_name": model_name,
        }
        LOADED_MODELS_CACHE[cache_key] = result
        return result

    @staticmethod
    def load_text_encoder(model_name: str, device: str = "cuda",
                          force_reload: bool = False):
        cache_key = f"{model_name}_{device}"
        if not force_reload and cache_key in LOADED_ENCODERS_CACHE:
            logger.info(f"Using cached text encoder: {model_name}")
            return LOADED_ENCODERS_CACHE[cache_key]

        from src.models.conditioner.qwen3_text_encoder import Qwen3TextEncoder

        config = TEXT_ENCODER_CONFIGS.get(model_name, {})
        embed_dim = config.get("embed_dim", 2048)
        max_length = config.get("max_length", 128)

        logger.info(f"Loading text encoder: {model_name}")
        conditioner = Qwen3TextEncoder(
            weight_path=model_name,
            embed_dim=embed_dim,
            max_length=max_length,
        )
        conditioner = conditioner.to(device)

        result = {
            "conditioner": conditioner,
            "device": device,
            "model_name": model_name,
        }
        LOADED_ENCODERS_CACHE[cache_key] = result
        return result

    @staticmethod
    def _resolve_checkpoint(repo_id: str, filename: str) -> str:
        """Download checkpoint from HF if not already local."""
        # Check ComfyUI models directory first
        import folder_paths
        try:
            local = folder_paths.get_full_path("pixelgen", filename)
            if local is not None and os.path.isfile(local):
                return local
        except Exception:
            pass

        logger.info(f"Downloading {filename} from {repo_id}...")
        path = hf_hub_download(repo_id=repo_id, filename=filename)
        return path