#!/usr/bin/env python3
"""
Replace ARM64 PIL binaries with x86_64 versions for Docker compatibility.

This script downloads the correct Pillow wheel for Python 3.10 x86_64
and replaces the current ARM64 binaries in src/lib-linux/.
"""
import urllib.request
import zipfile
import shutil
import os
from pathlib import Path

WHEEL_PATH = Path(__file__).parent.parent / "pillow-11.3.0-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl"
TEMP_EXTRACT = "/tmp/pillow_extracted"
TARGET_DIR = Path(__file__).parent.parent / "src" / "lib-linux"

def main():
    print("🔧 Replacing Linux PIL binaries: ARM64 → x86_64")
    
    # Step 1: Verify wheel exists
    print(f"\n📦 Using wheel: {WHEEL_PATH}")
    if not WHEEL_PATH.exists():
        print(f"❌ ERROR: Wheel not found at {WHEEL_PATH}")
        print("   Run: pip download --no-deps --only-binary=:all: --platform manylinux_2_28_x86_64 --python-version 310 --implementation cp --abi cp310 pillow")
        return
    print("✅ Wheel found")
    
    # Step 2: Extract wheel
    print(f"\n📦 Extracting wheel to: {TEMP_EXTRACT}")
    os.makedirs(TEMP_EXTRACT, exist_ok=True)
    with zipfile.ZipFile(WHEEL_PATH, 'r') as zip_ref:
        zip_ref.extractall(TEMP_EXTRACT)
    print("✅ Extraction complete")
    
    # Step 3: Remove old ARM64 binaries
    print(f"\n🗑️  Removing ARM64 binaries from: {TARGET_DIR}")
    if (TARGET_DIR / "PIL").exists():
        shutil.rmtree(TARGET_DIR / "PIL")
    if (TARGET_DIR / "pillow.libs").exists():
        shutil.rmtree(TARGET_DIR / "pillow.libs")
    print("✅ Old binaries removed")
    
    # Step 4: Copy x86_64 binaries
    print(f"\n📂 Copying x86_64 binaries to: {TARGET_DIR}")
    shutil.copytree(f"{TEMP_EXTRACT}/PIL", TARGET_DIR / "PIL")
    shutil.copytree(f"{TEMP_EXTRACT}/pillow.libs", TARGET_DIR / "pillow.libs")
    print("✅ New binaries installed")
    
    # Step 5: Verify architecture
    print("\n🔍 Verifying architecture...")
    so_files = list((TARGET_DIR / "PIL").glob("*.so"))
    if so_files:
        example = so_files[0].name
        if "x86_64" in example:
            print(f"✅ Architecture verified: {example}")
            print("\n🎉 SUCCESS! Linux PIL binaries are now x86_64 compatible")
            print("\n▶️  Run test: python -m pytest tests/e2e/test_pil_dependency_regression.py -q --maxfail=1 --tb=short -v")
        else:
            print(f"⚠️  WARNING: Expected x86_64 but found: {example}")
    else:
        print("❌ ERROR: No .so files found!")
    
    # Cleanup
    print("\n🧹 Cleaning up temporary files...")
    shutil.rmtree(TEMP_EXTRACT)
    print("✅ Cleanup complete")

if __name__ == "__main__":
    main()