from flask import Flask, request, send_file, jsonify
from flask_cors import CORS
from rembg import remove
from PIL import Image
import io
import os

app = Flask(__name__)
CORS(app, resources={r"/api/*": {"origins": "*"}})

@app.route('/')
def index():
    return jsonify({"message": "Background removal server is running!"})

@app.route('/api/remove-background', methods=['POST'])
def remove_background_api():
    if 'file' not in request.files:
        return jsonify({"error": "No file part in the request"}), 400
    file = request.files['file']
    if file.filename == '':
        return jsonify({"error": "No file selected for uploading"}), 400
    
    input_bytes = file.read()
    try:
        input_image = Image.open(io.BytesIO(input_bytes)).convert("RGBA")
        
        output_image = remove(
            input_image,
            model='u2net',  # ודא שאנחנו עדיין משתמשים במודל האיכותי
            alpha_matting=True,
            alpha_matting_foreground_threshold=240,
            alpha_matting_background_threshold=10,
            alpha_matting_erode_size=10
        )

        output_bytes_stream = io.BytesIO()
        output_image.save(output_bytes_stream, format="PNG")
        output_bytes = output_bytes_stream.getvalue()
        
        return send_file(
            io.BytesIO(output_bytes),
            mimetype='image/png',
            as_attachment=False,
            download_name='processed.png'
        )
    except Exception as e:
        print(f"Error processing image: {e}")
        return jsonify({"error": f"Failed to process image: {str(e)}"}), 500

if __name__ == '__main__':
    port = int(os.environ.get('PORT', 7860))
    app.run(host='0.0.0.0', port=port)