"""
ComfyUI Tensor Prism Node Pack
===============================

Advanced model merging and enhancement nodes for ComfyUI, providing sophisticated
techniques for blending, enhancing, and manipulating Stable Diffusion models with
GPU-optimized memory management.

Author: Arctenox
Version: 1.1.0
License: GPL-3.0
Repository: https://github.com/AstrionX/ComfyUI-Tensor-Prism-Node-Pack

Features:
- Advanced model merging with multiple interpolation methods
- Spectral frequency-domain merging
- Granular SDXL block control
- Sophisticated masking system
- GPU-optimized memory management
- Cross-platform compatibility (CUDA/MPS/CPU)
"""

import os
import sys
import importlib.util
from pathlib import Path

# Add current directory to path for imports
current_dir = Path(__file__).parent
sys.path.insert(0, str(current_dir))

try:
    from .TensorPrism_MainMerge import TensorPrism_MainMerge
except ImportError:
    TensorPrism_MainMerge = None

try:
    from .TensorPrism_SDXLBlockMerge import SDXLBlockMergeTensorPrism
except ImportError:
    SDXLBlockMergeTensorPrism = None

try:
    from .TensorPrism_SDXLAdvancedBlockmerge import SDXLAdvancedBlockMergeTensorPrism
except ImportError:
    SDXLAdvancedBlockMergeTensorPrism = None

try:
    from .TensorPrism_ModelMaskGenerator import TensorPrism_ModelMaskGenerator
except ImportError:
    TensorPrism_ModelMaskGenerator = None

try:
    from .TensorPrism_ModelKeyFilter import TensorPrism_ModelKeyFilter
except ImportError:
    TensorPrism_ModelKeyFilter = None

try:
    from .TensorPrism_ModelMaskBlender import TensorPrism_ModelMaskBlender
except ImportError:
    TensorPrism_ModelMaskBlender = None

try:
    from .TensorPrism_ModelWeightModifier import TensorPrism_ModelWeightModifier
except ImportError:
    TensorPrism_ModelWeightModifier = None

try:
    from .TensorPrism_WeightedTensorMerge import TensorPrism_WeightedMaskMerge as TensorPrism_WeightedTensorMerge_Class
except ImportError:
    TensorPrism_WeightedTensorMerge_Class = None

try:
    from .TensorPrism_Enhancer import ModelEnhancerTensorPrism
except ImportError:
    ModelEnhancerTensorPrism = None

try:
    from .TensorPrism_LayeredBlend import TensorPrism_LayeredBlend
except ImportError:
    TensorPrism_LayeredBlend = None

try:
    from .TensorPrism_Prism import TensorPrism_FastPrism
except ImportError:
    TensorPrism_FastPrism = None

# Version info
__version__ = "1.1.0"
__author__ = "Arctenox"
__description__ = "Advanced model merging and enhancement nodes for ComfyUI"

NODE_CLASS_MAPPINGS = {}

# Core merging nodes
if TensorPrism_MainMerge:
    NODE_CLASS_MAPPINGS["TensorPrism_MainMerge"] = TensorPrism_MainMerge

# SDXL specific nodes
if SDXLBlockMergeTensorPrism:
    NODE_CLASS_MAPPINGS["SDXL Block Merge (Tensor Prism)"] = SDXLBlockMergeTensorPrism

if SDXLAdvancedBlockMergeTensorPrism:
    NODE_CLASS_MAPPINGS["SDXLAdvancedBlockMergeTensorPrism"] = SDXLAdvancedBlockMergeTensorPrism

# Masking and filtering nodes
if TensorPrism_ModelMaskGenerator:
    NODE_CLASS_MAPPINGS["TensorPrism_ModelMaskGenerator"] = TensorPrism_ModelMaskGenerator

if TensorPrism_ModelKeyFilter:
    NODE_CLASS_MAPPINGS["TensorPrism_ModelKeyFilter"] = TensorPrism_ModelKeyFilter

if TensorPrism_ModelMaskBlender:
    NODE_CLASS_MAPPINGS["TensorPrism_ModelMaskBlender"] = TensorPrism_ModelMaskBlender

# Enhancement and transformation nodes
if TensorPrism_ModelWeightModifier:
    NODE_CLASS_MAPPINGS["TensorPrism_ModelWeightModifier"] = TensorPrism_ModelWeightModifier

if TensorPrism_WeightedTensorMerge_Class:
    NODE_CLASS_MAPPINGS["TensorPrism_WeightedTensorMerge"] = TensorPrism_WeightedTensorMerge_Class

if ModelEnhancerTensorPrism:
    NODE_CLASS_MAPPINGS["ModelEnhancerTensorPrism"] = ModelEnhancerTensorPrism

if TensorPrism_LayeredBlend:
    NODE_CLASS_MAPPINGS["TensorPrism_LayeredBlend"] = TensorPrism_LayeredBlend

if TensorPrism_FastPrism:
    NODE_CLASS_MAPPINGS["TensorPrism_Prism"] = TensorPrism_FastPrism

NODE_DISPLAY_NAME_MAPPINGS = {}

# Core merging nodes
if "TensorPrism_MainMerge" in NODE_CLASS_MAPPINGS:
    NODE_DISPLAY_NAME_MAPPINGS["TensorPrism_MainMerge"] = "Main Merge (Tensor Prism)"

# SDXL specific nodes
if "SDXL Block Merge (Tensor Prism)" in NODE_CLASS_MAPPINGS:
    NODE_DISPLAY_NAME_MAPPINGS["SDXL Block Merge (Tensor Prism)"] = "SDXL Block Merge (Tensor Prism)"

if "SDXLAdvancedBlockMergeTensorPrism" in NODE_CLASS_MAPPINGS:
    NODE_DISPLAY_NAME_MAPPINGS["SDXLAdvancedBlockMergeTensorPrism"] = "SDXL Advanced Block Merge (Tensor Prism)"

# Masking and filtering nodes
if "TensorPrism_ModelMaskGenerator" in NODE_CLASS_MAPPINGS:
    NODE_DISPLAY_NAME_MAPPINGS["TensorPrism_ModelMaskGenerator"] = "Model Mask Generator (Tensor Prism)"

if "TensorPrism_ModelKeyFilter" in NODE_CLASS_MAPPINGS:
    NODE_DISPLAY_NAME_MAPPINGS["TensorPrism_ModelKeyFilter"] = "Model Key Filter (Tensor Prism)"

if "TensorPrism_ModelMaskBlender" in NODE_CLASS_MAPPINGS:
    NODE_DISPLAY_NAME_MAPPINGS["TensorPrism_ModelMaskBlender"] = "Mask Blender (Tensor Prism)"

# Enhancement and transformation nodes
if "TensorPrism_ModelWeightModifier" in NODE_CLASS_MAPPINGS:
    NODE_DISPLAY_NAME_MAPPINGS["TensorPrism_ModelWeightModifier"] = "Model Weight Modifier (Tensor Prism)"

if "TensorPrism_WeightedTensorMerge" in NODE_CLASS_MAPPINGS:
    NODE_DISPLAY_NAME_MAPPINGS["TensorPrism_WeightedTensorMerge"] = "Weighted Tensor Merge (Tensor Prism)"

if "ModelEnhancerTensorPrism" in NODE_CLASS_MAPPINGS:
    NODE_DISPLAY_NAME_MAPPINGS["ModelEnhancerTensorPrism"] = "Model Enhancer (Tensor Prism)"

if "TensorPrism_LayeredBlend" in NODE_CLASS_MAPPINGS:
    NODE_DISPLAY_NAME_MAPPINGS["TensorPrism_LayeredBlend"] = "Layered Blend (Tensor Prism)"

if "TensorPrism_Prism" in NODE_CLASS_MAPPINGS:
    NODE_DISPLAY_NAME_MAPPINGS["TensorPrism_Prism"] = "Prism (Tensor Prism)"

NODE_CATEGORIES = {}

for node_key in NODE_CLASS_MAPPINGS.keys():
    if "MainMerge" in node_key or "LayeredBlend" in node_key:
        NODE_CATEGORIES[node_key] = "Tensor Prism/Core"
    elif "SDXL" in node_key:
        NODE_CATEGORIES[node_key] = "Tensor_Prism/Merge"
    elif "Mask" in node_key:
        NODE_CATEGORIES[node_key] = "Tensor_Prism/Mask"
    elif "Weight" in node_key or "Enhancer" in node_key:
        NODE_CATEGORIES[node_key] = "Tensor_Prism/Transform"
    else:
        NODE_CATEGORIES[node_key] = "Tensor_Prism/Core"

def check_dependencies():
    """Check if required dependencies are available."""
    required_packages = [
        ("torch", "PyTorch >= 1.12.0"),
        ("numpy", "NumPy >= 1.21.0"), 
        ("psutil", "psutil >= 5.8.0"),
    ]
    
    missing_packages = []
    
    for package_name, description in required_packages:
        try:
            importlib.import_module(package_name)
        except ImportError:
            missing_packages.append(description)
    
    if missing_packages:
        print(f"[TensorPrism] Warning: Missing dependencies:")
        for pkg in missing_packages:
            print(f"  - {pkg}")
        print("[TensorPrism] Some features may not work properly.")
    
    return len(missing_packages) == 0

def get_system_info():
    """Get system information for optimization."""
    try:
        import torch
        import psutil
        
        info = {
            "torch_version": torch.__version__,
            "cuda_available": torch.cuda.is_available(),
            "mps_available": hasattr(torch.backends, 'mps') and torch.backends.mps.is_available(),
            "cpu_count": psutil.cpu_count(),
            "total_ram": round(psutil.virtual_memory().total / (1024**3), 1),
        }
        
        if info["cuda_available"]:
            info["cuda_device_count"] = torch.cuda.device_count()
            info["cuda_memory"] = round(torch.cuda.get_device_properties(0).total_memory / (1024**3), 1)
        
        return info
    except Exception as e:
        print(f"[TensorPrism] Could not get system info: {e}")
        return {}

def print_welcome_message():
    """Print welcome message with system information."""
    print("\n" + "="*60)
    print("🎭 TensorPrism Node Pack Loaded")
    print("="*60)
    print(f"Version: {__version__}")
    print(f"Author: {__author__}")
    print(f"Nodes: {len(NODE_CLASS_MAPPINGS)}")
    
    # System info
    sys_info = get_system_info()
    if sys_info:
        print(f"\n📊 System Information:")
        print(f"  PyTorch: {sys_info.get('torch_version', 'Unknown')}")
        print(f"  CPU Cores: {sys_info.get('cpu_count', 'Unknown')}")
        print(f"  RAM: {sys_info.get('total_ram', 'Unknown')}GB")
        
        if sys_info.get("cuda_available"):
            print(f"  CUDA: Available ({sys_info.get('cuda_device_count', 0)} device(s))")
            print(f"  GPU Memory: {sys_info.get('cuda_memory', 'Unknown')}GB")
        elif sys_info.get("mps_available"):
            print(f"  MPS: Available (Apple Silicon)")
        else:
            print(f"  GPU: CPU fallback mode")
    
    print(f"\n🚀 Available Nodes:")
    for category in set(NODE_CATEGORIES.values()):
        print(f"\n  {category}:")
        for node_class, node_category in NODE_CATEGORIES.items():
            if node_category == category:
                display_name = NODE_DISPLAY_NAME_MAPPINGS.get(node_class, node_class)
                print(f"    • {display_name}")
    
    print(f"\n💡 Memory Management:")
    if sys_info.get("cuda_memory", 0) >= 24:
        print("  Recommended: Default settings (24GB+ GPU)")
    elif sys_info.get("cuda_memory", 0) >= 12:
        print("  Recommended: Memory limit 8GB, auto precision")
    elif sys_info.get("cuda_memory", 0) >= 8:
        print("  Recommended: Memory limit 6GB, CPU fallback for large merges")
    else:
        print("  Recommended: CPU processing for best stability")
    
    print("="*60 + "\n")

# Initialize the package
def __init_package():
    """Initialize the package and perform startup checks."""
    try:
        # Check dependencies
        deps_ok = check_dependencies()
        
        # Print welcome message
        print_welcome_message()
        
        if not deps_ok:
            print("[TensorPrism] Warning: Some dependencies are missing. Please install them for full functionality.")
        
        print("[TensorPrism] Package initialized successfully!")
        
    except Exception as e:
        print(f"[TensorPrism] Error during initialization: {e}")
        print("[TensorPrism] Package may not function correctly.")

# Run initialization
__init_package()

# Export for ComfyUI
__all__ = [
    "NODE_CLASS_MAPPINGS",
    "NODE_DISPLAY_NAME_MAPPINGS", 
    "__version__",
    "__author__",
    "__description__"
]