#!/usr/bin/env python3
"""
Setup script for CAI package installation in ColiFormer.
This script handles the installation of the CAI package with proper build flags
to avoid wheel naming issues that occur with standard pip install.
"""

import subprocess
import sys
import os
import importlib.util

def check_cai_installed():
    """Check if CAI package is already installed and working."""
    try:
        spec = importlib.util.find_spec("CAI")
        if spec is None:
            return False

        # Try to import the specific functions we need
        import CAI
        from CAI import CAI as cai_func, relative_adaptiveness
        print("✅ CAI package is already installed and working")
        return True
    except ImportError as e:
        print(f"❌ CAI package not found or not working: {e}")
        return False

def install_cai():
    """Install CAI package with proper build configuration."""
    print("🔧 Installing CAI package...")

    cai_repo = "git+https://github.com/Benjamin-Lee/CodonAdaptationIndex.git@b6e017a92c58829f6a5aec8c26a21262bc2a6610"

    try:
        # First ensure we have build tools
        print("Installing build dependencies...")
        subprocess.run([
            sys.executable, "-m", "pip", "install", "--upgrade",
            "setuptools>=65.0", "wheel>=0.37.0", "pip>=21.0"
        ], check=True, capture_output=True)

        # Try installing with --no-use-pep517 flag first (preferred method)
        print("Attempting CAI installation with legacy build system...")
        try:
            result = subprocess.run([
                sys.executable, "-m", "pip", "install",
                "--no-use-pep517", "--no-cache-dir", cai_repo
            ], check=True, capture_output=True, text=True)
            print("✅ CAI installed successfully with legacy build")
            return True
        except subprocess.CalledProcessError:
            print("⚠️ Legacy build failed, trying standard installation...")

        # Fallback to standard installation
        result = subprocess.run([
            sys.executable, "-m", "pip", "install", "--no-cache-dir", cai_repo
        ], check=True, capture_output=True, text=True)

        print("✅ CAI installed successfully with standard build")
        return True

    except subprocess.CalledProcessError as e:
        print(f"❌ CAI installation failed: {e}")
        if hasattr(e, 'stderr') and e.stderr:
            print(f"Error output: {e.stderr}")
        return False
    except Exception as e:
        print(f"❌ Unexpected error during CAI installation: {e}")
        return False

def verify_cai_installation():
    """Verify that CAI package is working correctly."""
    try:
        import CAI
        from CAI import CAI as cai_func, relative_adaptiveness

        # Test basic functionality
        test_sequences = ["ATGAAATAA", "ATGGGCTAA"]
        weights = relative_adaptiveness(sequences=test_sequences)
        cai_score = cai_func("ATGAAATAA", weights=weights)

        print(f"✅ CAI verification successful (test score: {cai_score:.3f})")
        return True

    except Exception as e:
        print(f"❌ CAI verification failed: {e}")
        return False

def main():
    """Main setup function."""
    print("ColiFormer CAI Setup")
    print("=" * 50)

    # Check if already installed
    if check_cai_installed():
        if verify_cai_installation():
            print("🎉 CAI is ready to use!")
            return True
        else:
            print("⚠️ CAI is installed but not working properly, reinstalling...")

    # Install CAI
    if install_cai():
        # Verify installation
        if verify_cai_installation():
            print("🎉 CAI setup completed successfully!")
            return True
        else:
            print("💥 CAI installation verification failed!")
            return False
    else:
        print("💥 CAI installation failed!")
        return False

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