#!/usr/bin/env python3
"""
Post-install script for Hugging Face Spaces
This script installs complex dependencies that need PyTorch to be available first
"""

import os
import sys
import subprocess
import shutil
from pathlib import Path

# Global variables for PyTorch version detection
PYTORCH_VERSION = None
CUDA_VERSION = None

def run_command(command, cwd=None, env=None):
    """Run a shell command and return the result"""
    print(f"Running: {command}")
    result = subprocess.run(command, shell=True, cwd=cwd, capture_output=True, text=True, env=env)
    if result.returncode != 0:
        print(f"Error running command: {command}")
        print(f"Error output: {result.stderr}")
        return False
    print(f"Success: {command}")
    return True

def check_pytorch_cuda():
    """Check if PyTorch is installed with correct CUDA version"""
    try:
        import torch
        print(f"PyTorch version: {torch.__version__}")
        print(f"CUDA available: {torch.cuda.is_available()}")
        if torch.cuda.is_available():
            print(f"CUDA version: {torch.version.cuda}")
            print(f"cuDNN version: {torch.backends.cudnn.version()}")
        return True
    except Exception as e:
        print(f"Error checking PyTorch: {e}")
        return False

def install_torch_sparse():
    """Install torch-sparse with compatible PyTorch version"""
    print("Installing torch-sparse...")
    
    # Check if torch-sparse is already installed
    try:
        import torch_sparse
        print("torch-sparse already installed")
        return True
    except ImportError:
        pass
    
    # Check current PyTorch version and adapt to it
    try:
        import torch
        print(f"Current PyTorch version: {torch.__version__}")
        
        # Extract PyTorch version info for compatibility
        if "+" in torch.__version__:
            pytorch_base = torch.__version__.split("+")[0]
            cuda_version = torch.__version__.split("+")[1]
        else:
            pytorch_base = torch.__version__
            cuda_version = "cpu"
        
        print(f"PyTorch base version: {pytorch_base}, CUDA: {cuda_version}")
        
        # Check if we have a very recent PyTorch version that might not have compatible wheels
        global PYTORCH_VERSION, CUDA_VERSION
        
        if pytorch_base.startswith("2.7"):
            print("PyTorch 2.7.x detected - this is very recent and may not have compatible wheels")
            print("Attempting to work with current PyTorch version and disable problematic extensions...")
            
            # Keep current PyTorch version but disable torch-sparse and torch-scatter
            PYTORCH_VERSION = pytorch_base
            CUDA_VERSION = cuda_version
            
            # Force disable torch-sparse and torch-scatter since they have ABI issues
            print("Disabling torch-sparse and torch-scatter due to ABI compatibility issues...")
            disable_torch_sparse()
            
            # Install torchvision for current PyTorch version
            print("Installing torchvision for current PyTorch version...")
            if not run_command("pip install torchvision"):
                print("Warning: Failed to install torchvision")
        else:
            # Store version info for later use
            PYTORCH_VERSION = pytorch_base
            CUDA_VERSION = cuda_version
        
    except ImportError:
        print("PyTorch not found, installing default version...")
        if not run_command("pip install torch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2 --index-url https://download.pytorch.org/whl/cu117"):
            return False
        PYTORCH_VERSION = "2.0.1"
        CUDA_VERSION = "cu117"
    
    # Check PyTorch installation
    print("Checking PyTorch installation...")
    check_pytorch_cuda()
    
    # Verify PyTorch version was actually updated
    try:
        import torch
        if not torch.__version__.startswith("2.0.1") or "+cu117" not in torch.__version__:
            print(f"Warning: PyTorch version is still {torch.__version__}, expected 2.0.1+cu117")
            print("This may cause compatibility issues with torch-sparse and torch-scatter")
    except Exception as e:
        print(f"Error checking PyTorch version: {e}")
    
    # Now install torch-sparse with the compatible version
    print(f"Installing torch-sparse with PyTorch {PYTORCH_VERSION}...")
    
    # Try to find compatible torch-sparse wheel
    wheel_url = f"https://data.pyg.org/whl/torch-{PYTORCH_VERSION}+{CUDA_VERSION}.html"
    print(f"Trying wheel URL: {wheel_url}")
    
    if run_command(f"pip install torch-sparse -f {wheel_url}"):
        print("Successfully installed torch-sparse")
        
        # Verify torch-sparse is compatible
        try:
            import torch_sparse
            print("torch-sparse import successful")
        except Exception as e:
            print(f"Warning: torch-sparse import failed: {e}")
            print("This may indicate a version compatibility issue")
        
        return True
    
    # If the specific wheel fails, try alternative versions
    print("Trying alternative torch-sparse versions...")
    alternative_versions = [
        f"https://data.pyg.org/whl/torch-{PYTORCH_VERSION}+{CUDA_VERSION}.html",
        f"https://data.pyg.org/whl/torch-{PYTORCH_VERSION}+cu118.html",
        f"https://data.pyg.org/whl/torch-{PYTORCH_VERSION}+cu117.html",
        f"https://data.pyg.org/whl/torch-{PYTORCH_VERSION}+cu116.html"
    ]
    
    for url in alternative_versions:
        if url != wheel_url:  # Skip the one we already tried
            print(f"Trying alternative URL: {url}")
            if run_command(f"pip install torch-sparse -f {url}"):
                print("Successfully installed torch-sparse with alternative version")
                return True
    
    # Try installing with --no-build-isolation
    print("Trying torch-sparse installation with --no-build-isolation...")
    if run_command(f"pip install torch-sparse -f {wheel_url} --no-build-isolation"):
        print("Successfully installed torch-sparse with --no-build-isolation")
        return True
    
    # Try installing from source with --no-build-isolation
    print("Trying torch-sparse installation from source...")
    if run_command("pip install torch-sparse --no-build-isolation"):
        print("Successfully installed torch-sparse from source")
        return True
    
    # Try installing with specific build flags for PyTorch 2.7
    print("Trying torch-sparse installation with specific build flags...")
    env = os.environ.copy()
    env['TORCH_CUDA_ARCH_LIST'] = '6.0;6.1;7.0;7.5;8.0;8.6'
    env['FORCE_CUDA'] = '1'
    if run_command("pip install torch-sparse --no-build-isolation", env=env):
        print("Successfully installed torch-sparse with build flags")
        return True
    
    # Try installing from git with specific version
    print("Trying torch-sparse installation from git...")
    if run_command("pip install git+https://github.com/rusty1s/pytorch_sparse.git --no-build-isolation"):
        print("Successfully installed torch-sparse from git")
        return True
    
    # If all else fails, disable torch-sparse and use built-in PyTorch sparse operations
    print("Failed to install torch-sparse, disabling it and using built-in PyTorch sparse operations")
    disable_torch_sparse()
    return True

def disable_torch_sparse():
    """Disable torch-sparse in the code by modifying the configuration"""
    try:
        # Modify the PoissonSystem.py to disable torch-sparse
        poisson_file = "NeuralJacobianFields/PoissonSystem.py"
        if os.path.exists(poisson_file):
            with open(poisson_file, 'r') as f:
                content = f.read()
            
            # Replace USE_TORCH_SPARSE = True with False
            content = content.replace("USE_TORCH_SPARSE = True", "USE_TORCH_SPARSE = False")
            
            with open(poisson_file, 'w') as f:
                f.write(content)
            
            print("Disabled torch-sparse in NeuralJacobianFields/PoissonSystem.py")
            print("Will use built-in PyTorch sparse operations instead")
    except Exception as e:
        print(f"Warning: Could not disable torch-sparse: {e}")

def install_torch_scatter():
    """Install torch-scatter with compatible PyTorch version"""
    print("Installing torch-scatter...")
    
    # Check if torch-scatter is already installed
    try:
        import torch_scatter
        print("torch-scatter already installed")
        return True
    except ImportError:
        pass
    
    # Install torch-scatter with the compatible PyTorch version
    print(f"Installing torch-scatter with PyTorch {PYTORCH_VERSION}...")
    
    # Try to find compatible torch-scatter wheel
    wheel_url = f"https://data.pyg.org/whl/torch-{PYTORCH_VERSION}+{CUDA_VERSION}.html"
    print(f"Trying wheel URL: {wheel_url}")
    
    if run_command(f"pip install torch-scatter -f {wheel_url}"):
        print("Successfully installed torch-scatter")
        
        # Verify torch-scatter is compatible
        try:
            import torch_scatter
            print("torch-scatter import successful")
        except Exception as e:
            print(f"Warning: torch-scatter import failed: {e}")
            print("This may indicate a version compatibility issue")
        
        return True
    
    # If the specific wheel fails, try alternative versions
    print("Trying alternative torch-scatter versions...")
    alternative_versions = [
        f"https://data.pyg.org/whl/torch-{PYTORCH_VERSION}+{CUDA_VERSION}.html",
        f"https://data.pyg.org/whl/torch-{PYTORCH_VERSION}+cu118.html",
        f"https://data.pyg.org/whl/torch-{PYTORCH_VERSION}+cu117.html",
        f"https://data.pyg.org/whl/torch-{PYTORCH_VERSION}+cu116.html"
    ]
    
    for url in alternative_versions:
        if url != wheel_url:  # Skip the one we already tried
            print(f"Trying alternative URL: {url}")
            if run_command(f"pip install torch-scatter -f {url}"):
                print("Successfully installed torch-scatter with alternative version")
                return True
    
    # Try installing with --no-build-isolation
    print("Trying torch-scatter installation with --no-build-isolation...")
    if run_command(f"pip install torch-scatter -f {wheel_url} --no-build-isolation"):
        print("Successfully installed torch-scatter with --no-build-isolation")
        return True
    
    # Try installing from source with --no-build-isolation
    print("Trying torch-scatter installation from source...")
    if run_command("pip install torch-scatter --no-build-isolation"):
        print("Successfully installed torch-scatter from source")
        return True
    
    # Try installing with specific build flags for PyTorch 2.7
    print("Trying torch-scatter installation with specific build flags...")
    env = os.environ.copy()
    env['TORCH_CUDA_ARCH_LIST'] = '6.0;6.1;7.0;7.5;8.0;8.6'
    env['FORCE_CUDA'] = '1'
    if run_command("pip install torch-scatter --no-build-isolation", env=env):
        print("Successfully installed torch-scatter with build flags")
        return True
    
    # Try installing from git with specific version
    print("Trying torch-scatter installation from git...")
    if run_command("pip install git+https://github.com/rusty1s/pytorch_scatter.git --no-build-isolation"):
        print("Successfully installed torch-scatter from git")
        return True
    
    # If all else fails, note that torch-scatter is not critical for basic functionality
    print("Failed to install torch-scatter, but this may not be critical for basic functionality")
    return True

def install_nvdiffrast():
    """Install nvdiffrast"""
    print("Installing nvdiffrast...")
    
    # Check if nvdiffrast is already installed
    try:
        import nvdiffrast
        print("nvdiffrast already installed")
        return True
    except ImportError:
        pass
    
    # Create packages directory if it doesn't exist
    packages_dir = Path("packages")
    packages_dir.mkdir(exist_ok=True)
    
    # Clone nvdiffrast
    if not (packages_dir / "nvdiffrast").exists():
        if not run_command("git clone https://github.com/NVlabs/nvdiffrast.git", cwd=packages_dir):
            return False
    
    # Install nvdiffrast
    nvdiffrast_dir = packages_dir / "nvdiffrast"
    if not run_command("pip install .", cwd=nvdiffrast_dir):
        return False
    
    return True

def install_pytorch3d():
    """Install PyTorch3D"""
    print("Installing PyTorch3D...")
    
    # Check if PyTorch3D is already installed
    try:
        import pytorch3d
        print(f"PyTorch3D already installed: {pytorch3d.__version__}")
        return True
    except ImportError:
        pass
    
    # Try installing from conda-forge (most reliable method)
    print("Trying PyTorch3D installation from conda-forge...")
    if run_command("pip install pytorch3d --index-url https://pypi.anaconda.org/conda-forge/simple/"):
        print("Successfully installed PyTorch3D from conda-forge")
        return True
    
    # Try the official PyTorch3D installation command for current CUDA version
    if PYTORCH_VERSION and CUDA_VERSION:
        # Try CUDA 12.6 wheels first (current system CUDA)
        if CUDA_VERSION == "cu126":
            print("Trying PyTorch3D installation for CUDA 12.6...")
            if run_command("pip install pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/py310_cu126_pyt271/download.html"):
                print("Successfully installed PyTorch3D for CUDA 12.6")
                return True
        # Try CUDA 11.7 wheels for older PyTorch versions
        elif CUDA_VERSION == "cu117":
            print("Trying PyTorch3D installation for CUDA 11.7...")
            if run_command("pip install pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/py310_cu117_pyt201/download.html"):
                print("Successfully installed PyTorch3D for CUDA 11.7")
                return True
    
    # Try to build PyTorch3D from source with system CUDA
    print("Trying PyTorch3D source installation with system CUDA...")
    if install_pytorch3d_from_source():
        return True
    
    # Fallback to CUDA 11.7 wheels
    print("Trying official PyTorch3D installation (CUDA 11.7)...")
    if run_command("pip install pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/py310_cu117_pyt201/download.html"):
        print("Successfully installed PyTorch3D from official wheels")
        return True
    
    # Try installing a specific version that's known to work with PyTorch 2.0.1
    print("Trying PyTorch3D installation with specific version...")
    if run_command("pip install pytorch3d==0.7.4"):
        print("Successfully installed PyTorch3D 0.7.4")
        return True
    
    # Try different wheel URLs for different Python/PyTorch versions
    # Use the detected PyTorch version
    print(f"Determining PyTorch3D wheel URL for PyTorch {PYTORCH_VERSION}+{CUDA_VERSION}")
    
    # Try specific wheel URL for current PyTorch version
    specific_url = f"https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/py310_{CUDA_VERSION}_pyt{PYTORCH_VERSION.replace('.', '')}/download.html"
    print(f"Trying specific wheel URL: {specific_url}")
    if run_command(f"pip install pytorch3d -f {specific_url}"):
        print(f"Successfully installed PyTorch3D with specific wheel URL")
        return True
    
    # Fallback to known working wheel URLs
    wheel_urls = [
        "https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/py310_cu126_pyt271/download.html",  # Current system CUDA
        f"https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/py310_{CUDA_VERSION}_pyt{PYTORCH_VERSION.replace('.', '')}/download.html",
        "https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/py310_cu117_pyt201/download.html",
        "https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/py39_cu117_pyt201/download.html",
        "https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/py38_cu117_pyt201/download.html"
    ]
    
    for i, url in enumerate(wheel_urls):
        print(f"Trying PyTorch3D wheel URL {i+1}...")
        if run_command(f"pip install pytorch3d -f {url}"):
            print(f"Successfully installed PyTorch3D from wheel URL {i+1}")
            return True
    
    # If all else fails, try source installation with proper build tools
    print("Trying PyTorch3D source installation with build tools...")
    packages_dir = Path("packages")
    
    # Clone PyTorch3D
    if not (packages_dir / "pytorch3d").exists():
        if not run_command("git clone https://github.com/facebookresearch/pytorch3d.git", cwd=packages_dir):
            return False
    
    # Install build dependencies
    print("Installing build dependencies...")
    run_command("pip install wheel setuptools ninja")
    
    # Install PyTorch3D with minimal setup
    pytorch3d_dir = packages_dir / "pytorch3d"
    
    # Try installation with build isolation disabled
    print("Trying PyTorch3D installation with build isolation disabled...")
    if run_command("pip install . --no-build-isolation", cwd=pytorch3d_dir):
        print("Successfully installed PyTorch3D from source")
        return True
    
    # Final fallback: try with environment variables
    print("Trying PyTorch3D installation with environment variables...")
    env = os.environ.copy()
    env['FORCE_CUDA'] = '1'
    env['CUDA_HOME'] = '/usr/local/cuda-11.7'
    env['CUDA_VERSION'] = '11.7'
    
    if run_command("pip install . --no-build-isolation", cwd=pytorch3d_dir, env=env):
        print("Successfully installed PyTorch3D with environment variables")
        return True
    
    return False

def install_pytorch3d_from_source():
    """Install PyTorch3D from source with system CUDA"""
    print("Installing PyTorch3D from source with system CUDA...")
    
    packages_dir = Path("packages")
    
    # Clone PyTorch3D
    if not (packages_dir / "pytorch3d").exists():
        if not run_command("git clone https://github.com/facebookresearch/pytorch3d.git", cwd=packages_dir):
            return False
    
    # Install build dependencies
    print("Installing build dependencies...")
    run_command("pip install wheel setuptools ninja")
    
    # Install PyTorch3D with system CUDA
    pytorch3d_dir = packages_dir / "pytorch3d"
    
    # Set environment variables to use system CUDA
    env = os.environ.copy()
    env['FORCE_CUDA'] = '1'
    env['CUDA_HOME'] = '/usr/local/cuda'
    env['CUDA_VERSION'] = '12.6'
    env['TORCH_CUDA_ARCH_LIST'] = '8.9'  # For L4 GPU
    
    print("Building PyTorch3D with system CUDA...")
    if run_command("pip install . --no-build-isolation", cwd=pytorch3d_dir, env=env):
        print("Successfully installed PyTorch3D from source with system CUDA")
        return True
    
    return False

def install_pytorch_dependencies():
    """Install PyTorch-related dependencies"""
    print("Installing PyTorch dependencies...")
    
    # Check if torchvision and torchaudio are already installed
    try:
        import torchvision
        import torchaudio
        print("PyTorch dependencies already installed")
        return True
    except ImportError:
        pass
    
    # Install torchvision if not present
    if not run_command("pip install torchvision"):
        print("Warning: Failed to install torchvision")
        return False
    
    # Install torchaudio if needed
    if not run_command("pip install torchaudio"):
        print("Warning: Failed to install torchaudio")
    
    return True

def install_fashion_clip():
    """Setup Fashion-CLIP"""
    print("Setting up Fashion-CLIP...")
    
    packages_dir = Path("packages")
    
    # Clone Fashion-CLIP if not already present
    if not (packages_dir / "fashion-clip").exists():
        if not run_command("git clone https://github.com/patrickjohncyh/fashion-clip.git", cwd=packages_dir):
            return False
    
    # Install Fashion-CLIP dependencies
    fashion_clip_dir = packages_dir / "fashion-clip"
    dependencies = ["appdirs", "boto3", "annoy", "validators", "transformers", "datasets"]
    
    for dep in dependencies:
        if not run_command(f"pip install {dep}", cwd=fashion_clip_dir):
            print(f"Warning: Failed to install {dep}")
    
    return True

def main():
    """Main installation function"""
    print("Starting post-installation for Garment3DGen...")
    
    # Install complex dependencies
    if not install_torch_sparse():
        print("Failed to install torch-sparse")
        sys.exit(1)
    
    if not install_torch_scatter():
        print("Failed to install torch-scatter")
        sys.exit(1)
    
    if not install_nvdiffrast():
        print("Failed to install nvdiffrast")
        sys.exit(1)
    
    if not install_pytorch3d():
        print("Failed to install PyTorch3D")
        sys.exit(1)
    
    if not install_fashion_clip():
        print("Failed to install Fashion-CLIP")
        sys.exit(1)
    
    if not install_pytorch_dependencies():
        print("Failed to install PyTorch dependencies")
        sys.exit(1)
    
    # Final verification
    print("\n=== Final Verification ===")
    print("Checking all dependencies...")
    
    try:
        import torch
        print(f"✓ PyTorch {torch.__version__} - CUDA: {torch.cuda.is_available()}")
        
        # Check if PyTorch version is compatible
        if PYTORCH_VERSION and not torch.__version__.startswith(PYTORCH_VERSION):
            print(f"⚠ Warning: PyTorch version {torch.__version__} may not be compatible with installed extensions")
            print(f"Expected version: {PYTORCH_VERSION}")
        
        # Check if torch-sparse is available or disabled
        try:
            import torch_sparse
            print("✓ torch-sparse")
        except ImportError:
            # Check if torch-sparse was disabled
            try:
                with open("NeuralJacobianFields/PoissonSystem.py", 'r') as f:
                    content = f.read()
                if "USE_TORCH_SPARSE = False" in content:
                    print("✓ torch-sparse (disabled, using built-in PyTorch sparse)")
                else:
                    print("✗ torch-sparse (not available)")
            except:
                print("✗ torch-sparse (not available)")
        except Exception as e:
            # Check if torch-sparse was disabled due to ABI issues
            try:
                with open("NeuralJacobianFields/PoissonSystem.py", 'r') as f:
                    content = f.read()
                if "USE_TORCH_SPARSE = False" in content:
                    print("✓ torch-sparse (disabled due to ABI issues, using built-in PyTorch sparse)")
                else:
                    print(f"✗ torch-sparse import failed: {e}")
            except:
                print(f"✗ torch-sparse import failed: {e}")
        
        # Check if torch-scatter is available
        try:
            import torch_scatter
            print("✓ torch-scatter")
        except ImportError:
            print("⚠ torch-scatter (not available, may not be critical)")
        except Exception as e:
            print(f"⚠ torch-scatter import failed: {e} (may not be critical)")
        
        import nvdiffrast
        print("✓ nvdiffrast")
        
        import pytorch3d
        print("✓ PyTorch3D")
        
        print("\nPost-installation completed successfully!")
        print("All dependencies are now available.")
        
    except ImportError as e:
        print(f"✗ Import error: {e}")
        print("This may indicate a version compatibility issue between PyTorch and its extensions")
        sys.exit(1)
    except Exception as e:
        print(f"✗ Verification error: {e}")
        print("This may indicate a version compatibility issue between PyTorch and its extensions")
        sys.exit(1)

if __name__ == "__main__":
    main()