"""
Background Removal API for Hugging Face Spaces
Self-hosted, FREE forever with unlimited usage!
"""
from flask import Flask, request, send_file
from flask_cors import CORS
from rembg import remove, new_session
from PIL import Image
import io
import os

app = Flask(__name__)

# Allow CORS from your Vercel domain
CORS(app, origins=["*"])  # In production, replace * with your Vercel domain

# Global session variable
_session = None

def get_session():
    """Lazy load rembg session on first use"""
    global _session
    if _session is None:
        print("🔧 Initializing rembg session with u2net model...")
        try:
            _session = new_session("u2net")
            print("✅ Session initialized successfully!")
        except Exception as e:
            print(f"❌ Session init error: {e}")
            raise
    return _session

@app.route('/')
def root():
    """Root endpoint with API info"""
    return {
        'name': 'Background Removal API',
        'version': '2.1.0',
        'backend': 'Self-hosted rembg (u2net model - High Quality)',
        'platform': 'Hugging Face Spaces',
        'features': '100% FREE forever, Unlimited usage, No API keys needed',
        'endpoints': {
            'POST /remove-background': 'Remove background from image',
            'GET /health': 'Health check'
        }
    }, 200

@app.route('/health')
def health_check():
    """Health check endpoint"""
    return {
        'status': 'ok', 
        'message': 'Background removal API is running',
        'backend': 'Self-hosted rembg with u2net model (High Quality)',
        'platform': 'Hugging Face Spaces (FREE forever)'
    }, 200

@app.route('/remove-background', methods=['POST'])
def remove_background():
    """
    Background removal using rembg with u2net model
    NO external APIs - 100% self-hosted!
    Supports quality options: standard, hd, ultra
    """
    try:
        # Get uploaded file
        if 'image' not in request.files:
            return {'error': 'No image file provided'}, 400
        
        file = request.files['image']
        input_image = Image.open(file.stream).convert('RGBA')
        
        # Get quality setting (default: ultra)
        quality = request.form.get('quality', 'ultra').lower()
        
        print(f"🎨 Processing image: {file.filename}")
        print(f"📐 Original size: {input_image.size}")
        print(f"🎯 Quality mode: {quality}")
        
        # Resize based on quality (preserve aspect ratio)
        max_dimensions = {
            'standard': 1920,  # Full HD
            'hd': 2560,        # 2K
            'ultra': 3840,     # 4K
            '4k': 3840         # 4K alias
        }
        
        max_size = max_dimensions.get(quality, 3840)  # Default to 4K
        
        # Only resize if image is larger than max
        if max(input_image.size) > max_size:
            ratio = max_size / max(input_image.size)
            new_size = tuple(int(dim * ratio) for dim in input_image.size)
            input_image = input_image.resize(new_size, Image.Resampling.LANCZOS)
            print(f"📐 Resized to: {input_image.size}")
        
        # Get session and remove background
        session = get_session()
        output_image = remove(input_image, session=session)
        
        print(f"✅ Background removed successfully!")
        
        # Save with MAXIMUM quality - no compression
        img_io = io.BytesIO()
        output_image.save(
            img_io, 
            'PNG',
            compress_level=0,  # No compression (0-9, 0=no compression)
            optimize=False     # Don't optimize file size
        )
        img_io.seek(0)
        
        output_size_mb = len(img_io.getvalue()) / (1024 * 1024)
        print(f"📦 Output size: {output_size_mb:.2f} MB")
        
        return send_file(
            img_io,
            mimetype='image/png',
            as_attachment=True,
            download_name=f"{file.filename.rsplit('.', 1)[0]}-nobg.png"
        )
    
    except Exception as e:
        print(f"❌ Error: {str(e)}")
        import traceback
        traceback.print_exc()
        return {'error': str(e)}, 500

if __name__ == '__main__':
    port = int(os.environ.get('PORT', 7860))  # HF Spaces uses port 7860
    print("🚀 Background Removal API starting on Hugging Face Spaces...")
    print(f"📍 Running on port {port}")
    print("✅ Using self-hosted rembg (NO external APIs!)")
    print("🎯 Model: u2netp (optimized for quality)")
    print("🆓 FREE forever - Unlimited usage!")
    app.run(host='0.0.0.0', port=port, debug=False)