"""ComfyUI-FL-PixelGen: Pixel-space diffusion text-to-image generation."""

import os
import sys
import logging

logger = logging.getLogger("FL_PixelGen")

# Add src/ to sys.path so bundled PixelGen code can be imported as `src.models.*`
_this_dir = os.path.dirname(os.path.abspath(__file__))
_src_parent = _this_dir  # src/ lives directly under this package
if _src_parent not in sys.path:
    sys.path.insert(0, _src_parent)

# Register model folder with ComfyUI
try:
    import folder_paths

    models_dir = os.path.join(
        folder_paths.models_dir, "pixelgen"
    )
    os.makedirs(models_dir, exist_ok=True)
    folder_paths.add_model_folder_path("pixelgen", models_dir)

    # Scan for user-placed checkpoints
    from modules.model_info import AVAILABLE_MODELS
    if os.path.isdir(models_dir):
        for f in os.listdir(models_dir):
            if f.endswith((".ckpt", ".pt", ".safetensors", ".pth")):
                name = os.path.splitext(f)[0]
                if name not in AVAILABLE_MODELS:
                    AVAILABLE_MODELS[name] = os.path.join(models_dir, f)
                    logger.info(f"Found local PixelGen model: {name}")
except Exception as e:
    logger.warning(f"Could not register PixelGen model folder: {e}")

# Import and register nodes
try:
    from .nodes import (
        FL_PixelGen_ModelLoader,
        FL_PixelGen_TextEncoderLoader,
        FL_PixelGen_Generate,
        FL_PixelGen_TrainingConfig,
        FL_PixelGen_Train,
        FL_PixelGen_LoraLoader,
    )

    NODE_CLASS_MAPPINGS = {
        "FL_PixelGen_ModelLoader": FL_PixelGen_ModelLoader,
        "FL_PixelGen_TextEncoderLoader": FL_PixelGen_TextEncoderLoader,
        "FL_PixelGen_Generate": FL_PixelGen_Generate,
        "FL_PixelGen_TrainingConfig": FL_PixelGen_TrainingConfig,
        "FL_PixelGen_Train": FL_PixelGen_Train,
        "FL_PixelGen_LoraLoader": FL_PixelGen_LoraLoader,
    }

    NODE_DISPLAY_NAME_MAPPINGS = {
        "FL_PixelGen_ModelLoader": "FL PixelGen Model Loader",
        "FL_PixelGen_TextEncoderLoader": "FL PixelGen Text Encoder Loader",
        "FL_PixelGen_Generate": "FL PixelGen Generate",
        "FL_PixelGen_TrainingConfig": "FL PixelGen Training Config",
        "FL_PixelGen_Train": "FL PixelGen Train",
        "FL_PixelGen_LoraLoader": "FL PixelGen LoRA Loader",
    }

    WEB_DIRECTORY = "./js"

    __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"]

except Exception as e:
    logger.error(f"Failed to load FL_PixelGen nodes: {e}")
    NODE_CLASS_MAPPINGS = {}
    NODE_DISPLAY_NAME_MAPPINGS = {}