"""
Build script to create PyArmor-obfuscated distributions for multiple Python versions.

This script builds separate wheels for Python 3.8, 3.9, and 3.10 for Windows and Linux.
"""
import os
import shutil
import subprocess
import sys
import platform
from pathlib import Path


# Package version
VERSION = "0.2.0"

# Python versions to build for
PYTHON_VERSIONS = ['3.8', '3.9', '3.10']

# Detect platform
IS_WINDOWS = platform.system() == 'Windows'
IS_LINUX = platform.system() == 'Linux'
PLATFORM_NAME = 'win' if IS_WINDOWS else 'linux' if IS_LINUX else 'unknown'


def clean_build_dirs():
    """Remove previous build artifacts."""
    dirs_to_clean = ['build', 'dist', 'orisys.egg-info', 'dist_obfuscated']
    for dir_name in dirs_to_clean:
        if os.path.exists(dir_name):
            print(f"Cleaning {dir_name}...")
            shutil.rmtree(dir_name)
    
    # Ensure dist directory exists with platform subdirectories
    os.makedirs(f'dist/{PLATFORM_NAME}', exist_ok=True)


def get_python_executable(version):
    """Get the Python executable for a specific version based on platform."""
    # First check if current Python matches the requested version
    current_version = f"{sys.version_info.major}.{sys.version_info.minor}"
    if current_version == version:
        return [sys.executable]
    
    if IS_WINDOWS:
        # Windows paths - including conda environments
        username = os.environ.get("USERNAME", "")
        possible_paths = [
            f'py -{version}',  # Windows py launcher
            f'python{version}',  # Direct version command
            # Conda environments with various naming patterns
            f'C:\\Users\\{username}\\anaconda3\\envs\\py{version.replace(".", "")}\\python.exe',
            f'C:\\Users\\{username}\\anaconda3\\envs\\py{version.replace(".", "")}-env\\python.exe',
            f'C:\\Users\\{username}\\anaconda3\\envs\\python{version.replace(".", "")}\\python.exe',
            f'C:\\Users\\{username}\\anaconda3\\envs\\orisys-py{version.replace(".", "")}\\python.exe',
            f'C:\\Users\\{username}\\miniconda3\\envs\\py{version.replace(".", "")}\\python.exe',
            # Standard Python installations
            f'C:\\Python{version.replace(".", "")}\\python.exe',
            f'C:\\Program Files\\Python{version.replace(".", "")}\\python.exe',
        ]
    else:
        # Linux/Unix paths
        possible_paths = [
            f'python{version}',
            f'/usr/bin/python{version}',
            f'/usr/local/bin/python{version}',
            f'~/.pyenv/versions/{version}/bin/python',
            f'/opt/python{version}/bin/python',
        ]
    
    for path in possible_paths:
        try:
            # Expand user home directory
            if path.startswith('~'):
                path = os.path.expanduser(path)
            
            if IS_WINDOWS and path.startswith('py -'):
                # Test Windows py launcher
                result = subprocess.run(
                    path.split() + ['-c', 'import sys; print(f"{sys.version_info.major}.{sys.version_info.minor}")'],
                    capture_output=True, text=True, shell=True
                )
                if result.returncode == 0 and result.stdout.strip() == version:
                    return path.split()
            else:
                # Check if file exists first
                if not os.path.exists(path):
                    continue
                    
                result = subprocess.run(
                    [path, '-c', 'import sys; print(f"{sys.version_info.major}.{sys.version_info.minor}")'],
                    capture_output=True, text=True
                )
                if result.returncode == 0 and result.stdout.strip() == version:
                    return [path]
        except (FileNotFoundError, subprocess.SubprocessError, PermissionError):
            continue
    
    return None


def obfuscate_code():
    """Obfuscate the source code using PyArmor (version-independent)."""
    print("\n=== Obfuscating code with PyArmor ===")
    
    # PyArmor 9.x generates version-independent obfuscated code
    # Output directly to dist_obfuscated/src (PyArmor will create orisys subdirectory)
    # Use --obf-code 1 to enable code obfuscation (default is 1, but explicit is better)
    # Use --exclude to exclude __init__.py files that don't need obfuscation (optional)
    if IS_WINDOWS:
        cmd = 'pyarmor gen --output dist_obfuscated/src --recursive --obf-code 1 src/orisys'
        print(f"Running: {cmd}")
        result = subprocess.run(cmd, capture_output=True, text=True, shell=True)
    else:
        cmd = [
            'pyarmor',
            'gen',
            '--output', 'dist_obfuscated/src',
            '--recursive',
            '--obf-code', '1',
            'src/orisys'
        ]
        print(f"Running: {' '.join(cmd)}")
        result = subprocess.run(cmd, capture_output=True, text=True)
    
    if result.returncode != 0:
        print(f"Error obfuscating code:\n{result.stderr}")
        print(f"stdout:\n{result.stdout}")
        sys.exit(1)
    
    print("Code obfuscation successful!")
    print(result.stdout)
    
    # Verify obfuscation: check that obfuscated files exist and original .py files are obfuscated
    obfuscated_dir = Path('dist_obfuscated/src/orisys')
    if obfuscated_dir.exists():
        py_files = list(obfuscated_dir.rglob('*.py'))
        print(f"Found {len(py_files)} obfuscated .py files")
        # Check a sample file to verify it's obfuscated (should contain PyArmor markers)
        if py_files:
            sample_file = py_files[0]
            try:
                with open(sample_file, 'rb') as f:
                    content = f.read(100)  # Read first 100 bytes
                    # PyArmor obfuscated files typically start with specific markers
                    if b'__pyarmor__' in content or b'pyarmor' in content.lower():
                        print(f"✓ Verified: {sample_file.name} is obfuscated")
                    else:
                        print(f"⚠ Warning: {sample_file.name} may not be properly obfuscated")
            except Exception as e:
                print(f"Could not verify obfuscation: {e}")
    
    return True


def copy_additional_files():
    """Copy non-Python files that don't need obfuscation."""
    print("\n=== Copying additional files ===")
    
    # Ensure configs directory exists in obfuscated output
    # IMPORTANT: Only copy JSON files, NOT .py files (those should be obfuscated by PyArmor)
    configs_src = Path('src/orisys/configs')
    configs_dst = Path('dist_obfuscated/src/orisys/configs')
    
    if configs_src.exists():
        print(f"Ensuring configs directory is copied (JSON files only)...")
        if not configs_dst.exists():
            configs_dst.mkdir(parents=True, exist_ok=True)
        
        # Only copy JSON files, not any .py files
        json_files_src = list(configs_src.glob('*.json'))
        json_files_dst = list(configs_dst.glob('*.json'))
        
        if len(json_files_dst) < len(json_files_src):
            print(f"Copying JSON files from configs directory ({len(json_files_src)} files)")
            for json_file in json_files_src:
                shutil.copy(json_file, configs_dst / json_file.name)
        else:
            print(f"Configs directory already present with {len(json_files_dst)} JSON files")
    
    # Copy config directory (if exists at root level)
    if os.path.exists('config'):
        print("Copying config directory...")
        shutil.copytree('config', 'dist_obfuscated/config', dirs_exist_ok=True)
    
    # Copy README and other metadata files
    files_to_copy = ['README.md', 'requirements.txt', 'LICENSE']
    for file in files_to_copy:
        if os.path.exists(file):
            print(f"Copying {file}...")
            shutil.copy(file, f'dist_obfuscated/{file}')
    
    # Verify that no original .py files are in dist_obfuscated (except obfuscated ones)
    print("\n=== Verifying obfuscated output ===")
    obfuscated_py_files = list(Path('dist_obfuscated/src/orisys').rglob('*.py'))
    print(f"Found {len(obfuscated_py_files)} .py files in obfuscated output")
    
    # Check if any files look like original source (not obfuscated)
    original_like_files = []
    for py_file in obfuscated_py_files:
        # Skip __init__.py files as they might be simple
        if py_file.name == '__init__.py':
            continue
        try:
            with open(py_file, 'rb') as f:
                content = f.read(500)  # Read first 500 bytes
                # Check for common Python source patterns that shouldn't be in obfuscated code
                if b'def ' in content and b'__pyarmor__' not in content:
                    # Might be original source
                    if b'import ' in content or b'from ' in content:
                        # Check if it has readable imports (obfuscated code has different structure)
                        text_content = content.decode('utf-8', errors='ignore')
                        if 'import ' in text_content and '__pyarmor__' not in text_content:
                            original_like_files.append(py_file)
        except Exception as e:
            print(f"Warning: Could not check {py_file}: {e}")
    
    if original_like_files:
        print(f"⚠️  Warning: Found {len(original_like_files)} files that may not be obfuscated:")
        for f in original_like_files[:5]:  # Show first 5
            print(f"   - {f}")
        if len(original_like_files) > 5:
            print(f"   ... and {len(original_like_files) - 5} more")
    else:
        print("✓ All .py files appear to be obfuscated")


def create_setup_for_obfuscated(python_version=None):
    """Create a modified setup.py for the obfuscated package."""
    print(f"\n=== Creating setup.py for obfuscated package (Python {python_version or '3.8-3.10'}) ===")
    
    if python_version:
        version_constraint = f'=={python_version}.*'
        py_tag = f'py{python_version.replace(".", "")}'
    else:
        version_constraint = '>=3.8,<3.11'
        py_tag = 'py38.py39.py310'
    
    setup_content = f'''from setuptools import setup, find_packages
import os
from pathlib import Path

# Read requirements
with open('requirements.txt', 'r') as f:
    requirements = [line.strip() for line in f if line.strip() and not line.startswith('#')]

# Find all packages including pyarmor_runtime
packages = find_packages(where="src")

# Collect all package data including PyArmor runtime files and configs
package_data = {{}}
for pkg in packages:
    pkg_path = Path('src') / pkg.replace('.', '/')
    if pkg_path.exists():
        data_files = []
        
        # Include all .pyd, .so, .dll files (PyArmor runtime)
        for pattern in ['*.pyd', '*.so', '*.dll', '*.dylib']:
            data_files.extend([str(f.relative_to(pkg_path)) for f in pkg_path.glob(pattern)])
        
        # Include JSON files from configs directory (for orisys.configs package)
        if pkg == 'orisys.configs':
            # Include all JSON files in the configs package directory
            for json_file in pkg_path.glob('*.json'):
                rel_path = json_file.relative_to(pkg_path)
                data_files.append(str(rel_path))
        
        # For orisys package, include configs subdirectory JSON files
        if pkg == 'orisys':
            configs_dir = pkg_path / 'configs'
            if configs_dir.exists() and configs_dir.is_dir():
                for json_file in configs_dir.glob('*.json'):
                    rel_path = json_file.relative_to(pkg_path)
                    data_files.append(str(rel_path))
        
        if data_files:
            package_data[pkg] = data_files

setup(
    name="orisys",
    version="{VERSION}",
    description="Orisys - Optical Tactile Sensor SDK (Protected)",
    author="Your Name",
    package_dir={{"": "src"}},
    packages=packages,
    package_data=package_data,
    include_package_data=True,
    install_requires=requirements,
    python_requires="{version_constraint}",
    # Include PyArmor runtime files
    zip_safe=False,
    # Exclude original source files - only include obfuscated files from dist_obfuscated
    exclude_package_data={{'': ['*.pyc', '*.pyo', '__pycache__']}},
    options={{
        'bdist_wheel': {{
            'python_tag': '{py_tag}',
        }}
    }},
)
'''
    
    with open('dist_obfuscated/setup.py', 'w') as f:
        f.write(setup_content)
    
    print("setup.py created successfully!")


def create_manifest_in():
    """Create MANIFEST.in to ensure JSON files in configs are included and exclude original source."""
    print("\n=== Creating MANIFEST.in ===")
    
    manifest_content = """# Include all JSON files in configs directory
include src/orisys/configs/*.json
recursive-include src/orisys/configs *.json

# Include PyArmor runtime files
recursive-include src *.pyd
recursive-include src *.so
recursive-include src *.dll
recursive-include src *.dylib

# Exclude original source files (only obfuscated .py files should be included)
# The obfuscated .py files from PyArmor will be included automatically
# But we want to make sure no original source files are accidentally included
prune src/orisys
recursive-include src/orisys *.py
"""
    
    manifest_path = Path('dist_obfuscated/MANIFEST.in')
    with open(manifest_path, 'w') as f:
        f.write(manifest_content)
    
    print("MANIFEST.in created successfully!")
    print(manifest_content)


def get_platform_tag():
    """Get the platform tag for wheel naming."""
    system = platform.system()
    machine = platform.machine().lower()
    
    if system == 'Windows':
        if machine in ('amd64', 'x86_64'):
            return 'win_amd64'
        elif machine in ('x86', 'i386', 'i686'):
            return 'win32'
        elif machine in ('arm64', 'aarch64'):
            return 'win_arm64'
    elif system == 'Linux':
        if machine in ('x86_64', 'amd64'):
            return 'manylinux2014_x86_64'
        elif machine in ('aarch64', 'arm64'):
            return 'manylinux2014_aarch64'
        elif machine in ('i386', 'i686'):
            return 'manylinux2014_i686'
    elif system == 'Darwin':
        if machine == 'arm64':
            return 'macosx_11_0_arm64'
        else:
            return 'macosx_10_9_x86_64'
    
    return 'any'


def build_wheel_for_version(python_version):
    """Build the wheel distribution for a specific Python version."""
    print(f"\n{'='*60}")
    print(f"Building wheel for Python {python_version} on {PLATFORM_NAME}")
    print('='*60)
    
    python_cmd = get_python_executable(python_version)
    
    if not python_cmd:
        print(f"⚠️  Python {python_version} not found on {PLATFORM_NAME}, skipping...")
        return False
    
    print(f"Using Python: {' '.join(python_cmd)}")
    
    # Verify Python version
    result = subprocess.run(
        python_cmd + ['-c', 'import sys; print(f"Python {sys.version}")'],
        capture_output=True, text=True
    )
    print(result.stdout.strip())
    
    # Get platform tag
    plat_tag = get_platform_tag()
    py_tag = f'py{python_version.replace(".", "")}'
    
    os.chdir('dist_obfuscated')
    
    # Build the wheel using setup.py with explicit tags
    cmd = python_cmd + ['setup.py', 'bdist_wheel', 
                        '--python-tag', py_tag,
                        '--plat-name', plat_tag,
                        '--dist-dir', f'../dist/{PLATFORM_NAME}']
    print(f"Running: {' '.join(cmd)}")
    result = subprocess.run(cmd, capture_output=True, text=True)
    
    os.chdir('..')
    
    if result.returncode != 0:
        print(f"❌ Error building wheel for Python {python_version}:\n{result.stderr}")
        return False
    
    print(f"✅ Wheel for Python {python_version} on {PLATFORM_NAME} built successfully!")
    print(f"   Expected: orisys-{VERSION}-{py_tag}-none-{plat_tag}.whl")
    return True


def build_universal_wheel():
    """Build a single universal wheel compatible with all Python versions."""
    print(f"\n{'='*60}")
    print(f"Building Universal Wheel (Python 3.8-3.10) for {PLATFORM_NAME}")
    print('='*60)
    
    os.chdir('dist_obfuscated')
    
    # Build the wheel
    cmd = [sys.executable, '-m', 'pip', 'wheel', '.', '--no-deps', '-w', f'../dist/{PLATFORM_NAME}']
    print(f"Running: {' '.join(cmd)}")
    result = subprocess.run(cmd, capture_output=True, text=True)
    
    os.chdir('..')
    
    if result.returncode != 0:
        print(f"❌ Error building universal wheel:\n{result.stderr}")
        return False
    
    print(f"✅ Universal wheel for {PLATFORM_NAME} built successfully!")
    print(result.stdout)
    return True


def main():
    """Main build process."""
    import argparse
    
    parser = argparse.ArgumentParser(description='Build orisys package with PyArmor')
    parser.add_argument('--mode', choices=['universal', 'multi-version', 'current'], default='current',
                       help='Build mode: universal (one wheel for all), multi-version (separate wheels), current (use current Python only)')
    parser.add_argument('--versions', nargs='+', default=PYTHON_VERSIONS,
                       help='Python versions to build for (only for multi-version mode)')
    
    args = parser.parse_args()
    
    # Get current Python version
    current_py_version = f"{sys.version_info.major}.{sys.version_info.minor}"
    
    print("=" * 60)
    print("Building Orisys Package with PyArmor Obfuscation")
    print(f"Platform: {PLATFORM_NAME} ({platform.system()})")
    print(f"Current Python: {current_py_version} ({sys.executable})")
    
    if args.mode == 'current':
        print(f"Mode: Building for current Python version ({current_py_version}) only")
        args.versions = [current_py_version]
        # Use multi-version logic but with current version only
        build_mode = 'multi-version'
    else:
        print(f"Python Versions: {', '.join(args.versions)}")
        print(f"Mode: {args.mode}")
        build_mode = args.mode
    
    print("=" * 60)
    
    # Ensure we're in the right directory
    script_dir = Path(__file__).parent
    os.chdir(script_dir)
    
    # Step 1: Clean previous builds
    clean_build_dirs()
    
    # Step 2: Obfuscate the code (once, version-independent)
    if not obfuscate_code():
        print("\n❌ Build failed during obfuscation!")
        sys.exit(1)
    
    # Step 3: Copy additional files (including configs)
    copy_additional_files()
    
    # Step 3.5: Create MANIFEST.in to ensure JSON files are included
    create_manifest_in()
    
    # Step 4 & 5: Build based on mode
    if build_mode == 'universal':
        # Build one wheel for all Python versions
        create_setup_for_obfuscated()
        if not build_universal_wheel():
            print("\n❌ Build failed!")
            sys.exit(1)
    else:
        # Build separate wheels for each Python version (or just current)
        success_count = 0
        for version in args.versions:
            create_setup_for_obfuscated(version)
            if build_wheel_for_version(version):
                success_count += 1
        
        if success_count == 0:
            print("\n❌ Build failed - no wheels created!")
            sys.exit(1)
        
        print(f"\n✅ Successfully built {success_count}/{len(args.versions)} wheels")
    
    print("\n" + "=" * 60)
    print("✅ Build completed successfully!")
    print("=" * 60)
    print(f"\nPlatform: {PLATFORM_NAME}")
    print(f"Obfuscated source: dist_obfuscated/")
    print(f"Distribution packages: dist/{PLATFORM_NAME}/")
    print("\nGenerated wheels:")
    for wheel in Path(f'dist/{PLATFORM_NAME}').glob('*.whl'):
        print(f"  - {wheel.name}")
    print("\nInstall with:")
    print(f"  pip install dist/{PLATFORM_NAME}/orisys-{VERSION}-*.whl")


if __name__ == '__main__':
    main()