#!/usr/bin/env python3
import subprocess
import sys
import platform
import os
import argparse

def run_command(command, check=True):
    """Run a shell command and print output"""
    print(f"Running: {command}")
    try:
        result = subprocess.run(
            command,
            shell=True,
            check=check,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True
        )
        print(result.stdout)
        if result.stderr:
            print(f"Error: {result.stderr}")
        return True
    except subprocess.CalledProcessError as e:
        print(f"Command failed with return code {e.returncode}")
        print(e.stderr)
        return False

def install_dependencies(use_conda=False, create_env=False, env_name="parser_env", upgrade=False, use_user_flag=False):
    """Install all required Python dependencies"""
    
    # Core dependencies
    dependencies = [
        "orjson",
        "tqdm",
        "matplotlib",
        "pyarrow",
        "numpy",
        "pandas",
        "polars"
    ]
    
    # Check if we're in a virtual environment
    in_venv = sys.prefix != sys.base_prefix
    in_conda = os.environ.get('CONDA_DEFAULT_ENV') is not None
    
    # Determine pip command
    pip_cmd = "pip"
    if platform.system() == "Windows":
        pip_cmd = "pip"  # on Windows, pip3 might not be recognized, use pip
    else:
        # Check if pip3 exists
        try:
            subprocess.run(["pip3", "--version"], check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            pip_cmd = "pip3"
        except (subprocess.SubprocessError, FileNotFoundError):
            pip_cmd = "pip"
    
    # Installation mode and flags
    install_mode = "--upgrade" if upgrade else "install"
    user_flag = "--user" if use_user_flag else ""
    
    if use_conda:
        # Conda installation
        if create_env:
            print(f"Creating new conda environment: {env_name}")
            if not run_command(f"conda create -y -n {env_name} python=3.9"):
                print("Failed to create conda environment")
                return False
            
            print(f"Activating conda environment: {env_name}")
            # We can't directly activate the environment in this script
            # Instead, we'll provide instructions for the user
            print("\nIMPORTANT: You need to activate the environment manually with:")
            print(f"conda activate {env_name}")
            print("Then run this script again without the --create-env option\n")
            
            return True
        
        # If using conda but not creating a new environment
        if not in_conda:
            print("Warning: You requested to use conda but you're not in a conda environment.")
            print("Please activate a conda environment first or use --create-env option.")
            return False
        
        # Install dependencies with conda when possible, fall back to pip
        for dep in dependencies:
            print(f"Installing {dep} with conda (will fall back to pip if needed)...")
            # Try conda first
            conda_success = run_command(f"conda install -y {dep}", check=False)
            if not conda_success:
                # Fall back to pip if conda fails
                print(f"Conda installation failed for {dep}, falling back to pip...")
                if not run_command(f"{pip_cmd} {install_mode} {dep}"):
                    print(f"Failed to install {dep}")
                    # Continue with other packages rather than stopping
    else:
        # Standard pip installation
        if not in_venv and not in_conda:
            print("Warning: You're not in a virtual environment.")
            print("It's recommended to use a virtual environment to avoid conflicts.")
            print("Consider using --conda or creating a venv manually.")
            response = input("Continue with system-wide installation? (y/n): ")
            if response.lower() not in ['y', 'yes']:
                print("Installation aborted.")
                return False
        
        # Install all dependencies at once
        deps_str = " ".join(dependencies)
        if not run_command(f"{pip_cmd} {install_mode} {user_flag} {deps_str}"):
            print("Failed to install dependencies")
            
            # Try installing one by one to see which ones succeed
            print("Attempting to install packages individually...")
            for dep in dependencies:
                run_command(f"{pip_cmd} {install_mode} {user_flag} {dep}", check=False)
    
    print("\nInstallation completed!")
    
    # Verify installations
    print("\nVerifying installations...")
    verification_failed = False
    for dep in dependencies:
        # Convert dependency name to import name (could be different)
        import_name = dep
        
        # Need to handle special cases
        if dep == "pyarrow":
            verification_code = "import pyarrow as pa; import pyarrow.dataset as ds; import pyarrow.ipc as ipc; print(f'PyArrow version: {pa.__version__}')"
        else:
            verification_code = f"import {import_name}; print(f'{import_name} version: ' + getattr({import_name}, '__version__', 'unknown'))"
        
        result = subprocess.run(
            [sys.executable, "-c", verification_code],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True
        )
        
        if result.returncode == 0:
            print(f"✓ {dep}: {result.stdout.strip()}")
        else:
            print(f"✗ {dep}: Failed to import")
            print(f"  Error: {result.stderr.strip()}")
            verification_failed = True
    
    if verification_failed:
        print("\nSome packages could not be imported. You may need to install them manually.")
    else:
        print("\nAll dependencies were successfully installed and verified!")
    
    return not verification_failed

def parse_args():
    parser = argparse.ArgumentParser(description='Install Python dependencies for the parser project')
    parser.add_argument('--conda', action='store_true', help='Use conda for installation when possible')
    parser.add_argument('--create-env', action='store_true', help='Create a new conda environment')
    parser.add_argument('--env-name', default='parser_env', help='Name for the conda environment (default: parser_env)')
    parser.add_argument('--upgrade', action='store_true', help='Upgrade existing packages')
    parser.add_argument('--create-venv', action='store_true', help='Create a Python virtual environment')
    parser.add_argument('--venv-name', default='venv', help='Name for the virtual environment (default: venv)')
    parser.add_argument('--brew', action='store_true', help='Show Homebrew installation commands for macOS')
    parser.add_argument('--user', action='store_true', help='Install packages in user space (with pip --user)')
    return parser.parse_args()

def create_venv(venv_dir="venv"):
    """Create a Python virtual environment"""
    if os.path.exists(venv_dir):
        print(f"Virtual environment '{venv_dir}' already exists.")
        return True
    
    print(f"Creating virtual environment in '{venv_dir}'...")
    result = subprocess.run(
        [sys.executable, "-m", "venv", venv_dir],
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        text=True
    )
    
    if result.returncode != 0:
        print(f"Failed to create virtual environment: {result.stderr}")
        return False
    
    # Determine the activate script path
    if platform.system() == "Windows":
        activate_script = os.path.join(venv_dir, "Scripts", "activate")
    else:
        activate_script = os.path.join(venv_dir, "bin", "activate")
    
    print("\nVirtual environment created successfully!")
    print(f"\nTo activate the environment, run:")
    if platform.system() == "Windows":
        print(f"{activate_script}")
    else:
        print(f"source {activate_script}")
    print("\nThen run this script again to install the packages in the virtual environment.")
    
    return True

def main():
    args = parse_args()
    
    # Check if we're on macOS - using multiple detection methods
    is_macos = False
    system_name = platform.system()
    if system_name == "Darwin":
        is_macos = True
    elif system_name == "":  # In some cases platform.system() might return empty
        # Try alternative detection methods
        try:
            # Check for macOS-specific directory
            is_macos = os.path.exists('/System/Library/CoreServices')
        except:
            pass
            
        if not is_macos:
            try:
                # Use uname as a fallback
                uname_output = subprocess.check_output(['uname', '-s'], text=True).strip()
                is_macos = uname_output == "Darwin"
            except:
                pass
    in_venv = sys.prefix != sys.base_prefix
    in_conda = os.environ.get('CONDA_DEFAULT_ENV') is not None
    
    # If on macOS and not in a virtual environment, handle the externally-managed environment
    if is_macos and not (in_venv or in_conda):
        print("Detected macOS with system Python.")
        print("On macOS, you must use a virtual environment to install packages (PEP 668).")
        print("Your options are:")
        print("1. Use a Python virtual environment (recommended)")
        print("2. Use Homebrew to install packages")
        print("3. Use Conda environment")
        
        choice = input("\nEnter your choice (1/2/3): ")
        
        if choice == "1":
            venv_name = input("Enter name for the virtual environment [venv]: ") or "venv"
            if create_venv(venv_name):
                return
            
        elif choice == "2":
            print("\nTo install with Homebrew, run these commands:")
            print("brew install python-orjson")
            print("brew install python-tqdm")
            print("brew install matplotlib")
            print("brew install apache-arrow") # This installs pyarrow
            print("brew install numpy")
            print("brew install pandas")
            print("brew install python-polars")
            print("\nNote: Some packages might need to be installed differently.")
            return
            
        elif choice == "3":
            if not shutil.which("conda"):
                print("Conda not found. Please install Miniconda or Anaconda first.")
                print("Visit: https://docs.conda.io/en/latest/miniconda.html")
                return
                
            args.conda = True
            args.create_env = True
    
    # Create a virtual environment if specified
    if hasattr(args, 'create_venv') and args.create_venv:
        venv_name = args.venv_name if hasattr(args, 'venv_name') else "venv"
        if create_venv(venv_name):
            return
    
    # Proceed with regular installation
    if install_dependencies(
        use_conda=args.conda,
        create_env=args.create_env,
        env_name=args.env_name,
        upgrade=args.upgrade,
        use_user_flag=args.user
    ):
        print("\nSetup completed successfully!")
    else:
        print("\nSetup completed with some issues. Please review the output above.")
        sys.exit(1)

if __name__ == "__main__":
    main()