#!/usr/bin/env python3
"""
Installation script for ComfyUI-HunyuanWorld
This script handles the installation of dependencies and setup.
"""

import subprocess
import sys
import os
import importlib.util

def check_package(package_name):
    """Check if a package is installed"""
    spec = importlib.util.find_spec(package_name)
    return spec is not None

def install_package(package):
    """Install a package using pip"""
    try:
        subprocess.check_call([sys.executable, "-m", "pip", "install", package])
        return True
    except subprocess.CalledProcessError as e:
        print(f"Failed to install {package}: {e}")
        return False

def install_requirements():
    """Install requirements from requirements.txt"""
    requirements_file = os.path.join(os.path.dirname(__file__), "requirements.txt")
    
    if not os.path.exists(requirements_file):
        print("❌ requirements.txt not found!")
        return False
    
    try:
        print("📦 Installing requirements...")
        subprocess.check_call([
            sys.executable, "-m", "pip", "install", "-r", requirements_file
        ])
        print("✅ Requirements installed successfully!")
        return True
    except subprocess.CalledProcessError as e:
        print(f"❌ Failed to install requirements: {e}")
        return False

def check_comfyui():
    """Check if we're in a ComfyUI environment"""
    # Look for ComfyUI indicators
    current_dir = os.getcwd()
    parent_dirs = []
    
    # Walk up the directory tree
    path = current_dir
    for _ in range(5):  # Check up to 5 levels up
        parent_dirs.append(path)
        parent_path = os.path.dirname(path)
        if parent_path == path:  # Reached root
            break
        path = parent_path
    
    # Check for ComfyUI indicators
    for dir_path in parent_dirs:
        if any(os.path.exists(os.path.join(dir_path, indicator)) for indicator in [
            "main.py", "nodes.py", "comfy", "custom_nodes"
        ]):
            return True
    
    return False

def main():
    """Main installation function"""
    print("🌍 ComfyUI-HunyuanWorld Installation")
    print("=" * 50)
    
    # Check if we're in ComfyUI environment
    if not check_comfyui():
        print("⚠️  Warning: This doesn't appear to be a ComfyUI installation directory.")
        print("   Make sure you're installing this in ComfyUI/custom_nodes/")
        response = input("   Continue anyway? (y/N): ")
        if response.lower() != 'y':
            print("Installation cancelled.")
            return False
    
    # Check Python version
    if sys.version_info < (3, 8):
        print("❌ Python 3.8 or higher is required!")
        return False
    
    print(f"✅ Python {sys.version_info.major}.{sys.version_info.minor} detected")
    
    # Install requirements
    if not install_requirements():
        return False
    
    # Check critical packages
    critical_packages = ["torch", "transformers", "diffusers", "trimesh"]
    missing_packages = []
    
    for package in critical_packages:
        if not check_package(package):
            missing_packages.append(package)
    
    if missing_packages:
        print(f"⚠️  Some critical packages may not be properly installed: {missing_packages}")
        print("   You may need to install them manually or restart your environment.")
    
    # Optional packages check
    optional_packages = {
        "hy3dgen": "Hunyuan3D library (for real pipeline functionality)",
        "gradio": "Gradio interface support",
        "fastapi": "API server support"
    }
    
    print("\n📋 Optional Package Status:")
    for package, description in optional_packages.items():
        status = "✅ Installed" if check_package(package) else "❌ Not installed"
        print(f"   {package}: {status} - {description}")
    
    print("\n🎉 Installation completed!")
    print("\n📚 Next Steps:")
    print("1. Restart ComfyUI")
    print("2. Look for 'HunyuanWorld' nodes in the node menu")
    print("3. Check the README.md for usage examples")
    print("4. If you encounter issues, check the troubleshooting section")
    
    # Hunyuan3D specific installation note
    if not check_package("hy3dgen"):
        print("\n💡 For full functionality, install Hunyuan3D:")
        print("   git clone https://github.com/tencent/hunyuan3d-2.git")
        print("   cd hunyuan3d-2")
        print("   pip install -e .")
    
    return True

if __name__ == "__main__":
    success = main()
    sys.exit(0 if success else 1)