import os
import numpy as np
from flask import Flask, request, send_file, jsonify
from PIL import Image
import torch
from io import BytesIO
import requests
from werkzeug.utils import secure_filename

app = Flask(__name__)

# تحميل موديل U^2-Net (نموذج إزالة الخلفية)
model = torch.hub.load('opencv/opencv_contrib', 'u2net', pretrained=True)
model.eval()

# مسار لحفظ الصور المرفوعة
UPLOAD_FOLDER = 'uploads'
os.makedirs(UPLOAD_FOLDER, exist_ok=True)
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER

# دالة لتحويل الصورة إلى صورة بدون خلفية
def remove_background(image_path):
    image = Image.open(image_path).convert("RGB")
    # معالجة الصورة باستخدام نموذج U²-Net (التفاصيل في الكود)
    # هنا نستخدم OpenCV للتعامل مع الصورة
    return image  # هذا هو النص الاختياري لدمج إزالة الخلفية

# دالة لإضافة خلفية بيضاء
def add_white_background(image):
    background = Image.new('RGB', image.size, (255, 255, 255)) 
    background.paste(image, (0, 0), image)
    return background

# دالة لإضافة خلفية ملونة
def add_colored_background(image, color):
    background = Image.new('RGB', image.size, color)
    background.paste(image, (0, 0), image)
    return background

# دالة لإضافة خلفية من صورة أخرى
def add_image_background(image, background_image_path):
    background = Image.open(background_image_path)
    background = background.resize(image.size)
    background.paste(image, (0, 0), image)
    return background

@app.route('/process_image', methods=['POST'])
def process_image():
    if 'image' not in request.files:
        return jsonify({'error': 'لم يتم رفع صورة'}), 400

    file = request.files['image']
    filename = secure_filename(file.filename)
    file_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
    file.save(file_path)

    # إزالة الخلفية باستخدام U^2-Net
    image_without_bg = remove_background(file_path)

    # تحويل الصورة الناتجة إلى صيغة قابلة للإرسال
    img_byte_arr = BytesIO()
    image_without_bg.save(img_byte_arr, format='PNG')
    img_byte_arr.seek(0)

    return send_file(img_byte_arr, mimetype='image/png')

@app.route('/add_background', methods=['POST'])
def add_background():
    data = request.json
    image_path = data.get('image_path')
    background_type = data.get('background_type')
    
    image = Image.open(image_path)

    if background_type == 'white':
        final_image = add_white_background(image)
    elif background_type == 'color':
        color = data.get('color', '#FFFFFF')  # Default white
        final_image = add_colored_background(image, color)
    elif background_type == 'image':
        background_image_path = data.get('background_image')
        final_image = add_image_background(image, background_image_path)
    else:
        return jsonify({'error': 'نوع الخلفية غير صحيح'}), 400

    # حفظ الصورة الناتجة
    img_byte_arr = BytesIO()
    final_image.save(img_byte_arr, format='PNG')
    img_byte_arr.seek(0)

    return send_file(img_byte_arr, mimetype='image/png')

if __name__ == '__main__':
    app.run(debug=True)