import gradio as gr
import numpy as np
import os
from PIL import Image
import requests
from io import BytesIO
import io
import base64

# استرداد مفتاح API من متغيرات البيئة
hf_token = os.environ.get("HF_TOKEN_API_DEMO")  # تأكد من ضبط هذا المتغير في بيئتك
auth_headers = {"Authorization": f"Bearer {hf_token}"}  # إعداد المصادقة

# تحويل الصورة إلى Base64
def convert_image_to_base64(image):
    buffer = io.BytesIO()
    image.save(buffer, format="PNG")  # يمكن استخدام تنسيقات أخرى مثل JPEG إذا لزم الأمر
    return base64.b64encode(buffer.getvalue()).decode('utf-8')

# تنزيل الصورة الناتجة من URL
def download_image(url):
    response = requests.get(url)
    return Image.open(BytesIO(response.content)).convert("RGB")

# استدعاء Bria API باستخدام نقطة النهاية gen_fill
def call_gen_fill_api(image, mask, api_token):
    url = "https://engine.prod.bria-api.com/v1/gen_fill"
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {api_token}",
    }
    
    # تحويل الصورة والقناع إلى Base64
    image_base64 = convert_image_to_base64(image)
    mask_base64 = convert_image_to_base64(mask)
    
    payload = {
        "file": image_base64,
        "mask_file": mask_base64,
        "mask_type": "manual",
        "prompt": "Erase object",
        "negative_prompt": "",
        "num_results": 1,
        "sync": True,  # معالجة متزامنة للحصول على النتيجة مباشرة
        "seed": 0,
        "content_moderation": False
    }
    
    response = requests.post(url, json=payload, headers=headers)
    if response.status_code == 200:
        result = response.json()
        result_url = result.get("result_url")
        if result_url:
            return download_image(result_url)  # تنزيل الصورة الناتجة
        else:
            raise Exception("No result URL found in the response.")
    else:
        raise Exception(f"API Error: {response.status_code} - {response.text}")

# دالة معالجة الصورة والقناع
def predict(dict):
    init_image = Image.fromarray(dict['background'][:, :, :3], 'RGB')
    mask = Image.fromarray(dict['layers'][0][:, :, 3], 'L')  # القناع بتدرج الرمادي
    
    api_token = hf_token  # استخدم مفتاح API الخاص بك
    if not api_token:
        return "Error: API token is missing. Please set it in your environment variables."
    
    try:
        result_image = call_gen_fill_api(init_image, mask, api_token)
        return result_image
    except Exception as e:
        return f"Error: {e}"

# تخصيص CSS
css = '''
.gradio-container{max-width: 1100px !important}
#image_upload{min-height:400px}
#image_upload [data-testid="image"], #image_upload [data-testid="image"] > div{min-height: 400px}
#run_button {
    width: 100%;
    height: 50px;
    display: flex;
    align-items: center;
    justify-content: center;
}
#output-img img, #image_upload img {
    object-fit: contain;
    width: 100%;
    height: auto;
}
'''

# إعداد واجهة Gradio
image_blocks = gr.Blocks(css=css, elem_id="total-container")
with image_blocks as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown("## BRIA Eraser API Integration")
        gr.HTML('''
          <p style="margin-bottom: 10px; font-size: 94%">
            This demo showcases the BRIA Gen-Fill API, allowing users to remove specific elements or objects from images.<br>
            Notes:<br>
            - For high-resolution images, processing time may be longer.<br>
            - Ensure masks are accurate for better results.<br>
          </p>
        ''')
    with gr.Row():
        with gr.Column():
            image = gr.ImageEditor(
                sources=["upload"],
                layers=False,
                transforms=[], 
                brush=gr.Brush(colors=["#000000"], color_mode="fixed"),  # اللون الأسود كخيار فرشاة
            )
            with gr.Row(equal_height=True):
                btn = gr.Button("Erase!", elem_id="run_button")
        
        with gr.Column():
            image_out = gr.Image(label="Output", elem_id="output-img")

    # تشغيل الزر لتفعيل وظيفة التعديل
    btn.click(fn=predict, inputs=[image], outputs=[image_out], api_name='run')
    
    gr.HTML(
        """
            <div class="footer">
                <p>Powered by BRIA API - Gradio Demo by Hugging Face</p>
            </div>
        """
    )

image_blocks.queue(max_size=25, api_open=False).launch(show_api=False)