#!/usr/bin/env python3
"""
Post-install script to patch numexpr to handle invalid OMP_NUM_THREADS values
"""
import os
import site

def patch_numexpr():
    """Find and patch numexpr/utils.py to handle invalid OMP_NUM_THREADS"""
    for site_dir in site.getsitepackages():
        numexpr_utils = os.path.join(site_dir, 'numexpr', 'utils.py')
        if os.path.exists(numexpr_utils):
            print(f"Found numexpr at: {numexpr_utils}")
            
            with open(numexpr_utils, 'r') as f:
                content = f.read()
            
            # Replace the problematic line
            old_code = "requested_threads = int(os.environ['OMP_NUM_THREADS'])"
            new_code = """try:
                requested_threads = int(os.environ['OMP_NUM_THREADS'])
            except (ValueError, TypeError):
                requested_threads = 4  # Default for HF Spaces T4 GPU"""
            
            if old_code in content and new_code not in content:
                content = content.replace(old_code, new_code)
                
                with open(numexpr_utils, 'w') as f:
                    f.write(content)
                
                print("✅ Patched numexpr successfully!")
                return True
            else:
                print("Already patched or couldn't find code to patch")
    
    return False

if __name__ == '__main__':
    patch_numexpr()