# image_agent_gradio/app.py  
# وكيل توليد الصور (محدث لدعم وكيل الفيديو)
# ============================================================  
  
import os  
import io  
import base64  
import logging  
from typing import Optional, Dict, Any, Tuple  
import requests  
import gradio as gr  
from PIL import Image  
from gradio_client import Client as GradioClient
  
# -------------------------- استيرادات الكممَة والبيئة -------------------------  
try:  
    from optimum.intel.openvino import OVDiffusionPipeline  
    import torch  
    LOCAL_OPTIMUM = True  
except ImportError:  
    LOCAL_OPTIMUM = False  
    try:  
        from diffusers import StableDiffusionPipeline  
        import torch  
    except ImportError:  
        pass   
  
logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")  
log = logging.getLogger("image_agent")  
  
if 'OVDiffusionPipeline' in globals():  
     log.info("✅ Optimum/OpenVINO متاح للتحسين على CPU باستخدام OVDiffusionPipeline.")  
elif 'StableDiffusionPipeline' in globals():  
     log.info("⚠️ OpenVINO غير متاح. سيتم استخدام Diffusers القياسي.")  
else:  
    log.warning("❌ لم يتم العثور على مكتبات Diffusers أو Optimum.")  
  
  
LOCAL_DIFFUSERS = LOCAL_OPTIMUM or ('StableDiffusionPipeline' in globals())  
  
HF_TOKEN = os.getenv("HF_API_TOKEN")  
HF_MODEL = os.getenv("HF_MODEL", "OpenVINO/stable-diffusion-xl-base-1.0-int8-ov")

# عنوان وكيل الفيديو
VIDEO_AGENT_URL = os.getenv("VIDEO_AGENT_URL", "https://mustafa-albakkar-videoagent.hf.space")
  
# ---------------- Initialization Check ----------------  
if LOCAL_OPTIMUM:  
    STATUS_MESSAGE = f"✅ الوضع المحلي (OpenVINO Optimized) نشط: {HF_MODEL}"  
elif LOCAL_DIFFUSERS:  
     STATUS_MESSAGE = f"⚠️ الوضع المحلي (Diffusers) نشط: {HF_MODEL} - قد يكون بطيئاً جداً على CPU."  
else:  
    if HF_TOKEN:  
        STATUS_MESSAGE = f"✅ وضع واجهة HF API نشط: {HF_MODEL}"  
    else:  
        STATUS_MESSAGE = "❌ خطأ: لم يتم تعيين HF_API_TOKEN، والتوليد المحلي غير نشط."  
  
  
# ---------------- Local Pipeline Cache ----------------  
_pipeline = None  
def get_local_pipeline():  
    """تهيئة وإرجاع خط أنابيب Diffusers/Optimum المحلي."""  
    global _pipeline  
    if _pipeline is None:  
        if not LOCAL_DIFFUSERS:  
             raise RuntimeError("التوليد المحلي غير متاح. تأكد من إعداد الاعتماديات.")  
          
        if HF_MODEL is None:  
            log.error("HF_MODEL is None. Cannot load pipeline.")  
            raise ValueError("❌ خطأ: لم يتم تعيين مسار النموذج (HF_MODEL).")  
          
        log.info(f"Loading local diffusers pipeline from: {HF_MODEL}")  
        device = "cuda" if torch.cuda.is_available() else "cpu"  
          
        pipeline_args = {}  
        if device == "cuda":  
            pipeline_args['torch_dtype'] = torch.float16  
          
        if LOCAL_OPTIMUM:  
            log.info("Loading OVDiffusionPipeline (Quantized pipeline).")  
            _pipeline = OVDiffusionPipeline.from_pretrained(HF_MODEL)  
              
        else:  
            log.info("Loading standard StableDiffusionPipeline.")  
            _pipeline = StableDiffusionPipeline.from_pretrained(  
                HF_MODEL,   
                **pipeline_args  
            )  
              
            _pipeline = _pipeline.to(device)  
          
    return _pipeline  
  
  
# ---------------- Inference Functions (Synchronous) ----------------  
  
def hf_inference(prompt: str) -> Tuple[Optional[Image.Image], str]:  
    """يولد صورة عبر واجهة Hugging Face Inference API."""  
    if not HF_TOKEN:  
        return None, "❌ فشل: لم يتم تعيين مفتاح 'HF_API_TOKEN'."  
      
    url = f"https://api-inference.huggingface.co/models/{HF_MODEL}"  
    headers = {"Authorization": f"Bearer {HF_TOKEN}"}  
    payload = {"inputs": prompt}  
      
    try:  
        log.info(f"Sending prompt to HF API: {prompt[:50]}...")  
        r = requests.post(url, headers=headers, json=payload, timeout=120)  
        r.raise_for_status()  
          
        if r.headers.get("content-type", "").startswith("image"):  
            image_bytes = r.content  
            image = Image.open(io.BytesIO(image_bytes))  
            return image, "✅ تم التوليد بنجاح (HF Inference API)."  
  
        data = r.json()  
          
        if isinstance(data, dict) and data.get("error"):  
            return None, f"❌ خطأ HF: {data.get('error')}"  
          
        return None, "❌ فشل: تنسيق استجابة غير مدعوم من HF API."  
  
    except requests.exceptions.HTTPError as e:  
        status_code = e.response.status_code  
        if status_code == 503:  
             return None, "⚠️ الخدمة غير متوفرة (503): قد تكون محملة أو غير جاهزة. حاول مرة أخرى."  
        return None, f"❌ خطأ HTTP: {status_code}. التفاصيل: {e}"  
    except Exception as e:  
        log.error("HF request failed: %s", e)  
        return None, f"❌ فشل الطلب: {e}"  
  
  
def local_generate(prompt: str) -> Tuple[Optional[Image.Image], str]:  
    """يولد صورة باستخدام خط أنابيب Diffusers/Optimum المحلي."""  
    try:  
        pipe = get_local_pipeline()  
          
        INFERENCE_STEPS = 70   
        log.info(f"Local generation started for prompt: {prompt[:50]}... with {INFERENCE_STEPS} steps.")  
          
        result = pipe(prompt, num_inference_steps=INFERENCE_STEPS)   
          
        if not hasattr(result, 'images') or result.images is None:  
            log.error("Pipeline result does not contain valid images property or it is None.")  
            return None, "❌ فشل التوليد المحلي: خط أنابيب التوليد لم يرجع أي قائمة صور صالحة."  
              
        valid_images = [img for img in result.images if img is not None]  
  
        if hasattr(result, 'nsfw_content_detected') and any(result.nsfw_content_detected):  
            return None, "❌ فشل التوليد المحلي: تم اكتشاف محتوى غير آمن."  
  
        if valid_images:  
            return valid_images[0], f"✅ تم التوليد بنجاح (محليًا) في {INFERENCE_STEPS} خطوة."  
        else:  
            return None, "❌ فشل التوليد المحلي: لم يتم إرجاع صور صالحة."  
              
    except Exception as e:  
        log.error("Local generation error: %s", e)  
        return None, f"❌ خطأ التوليد المحلي: {type(e).__name__}: {e}"  


# ---------------- NEW: Video Agent Integration ----------------

def send_to_video_agent(
    image_base64: str,
    quote: str,
    author: str,
    culture: str
) -> Dict[str, Any]:
    """
    إرسال البيانات إلى وكيل الفيديو للمعالجة.
    
    Returns:
        Dict containing video_base64 and status
    """
    try:
        log.info(f"Sending data to Video Agent at {VIDEO_AGENT_URL}...")
        
        client = GradioClient(VIDEO_AGENT_URL)
        
        result = client.predict(
            image_base64,
            quote,
            author,
            culture,
            api_name="/generate_video"
        )
        
        log.info("✅ Video Agent response received successfully.")
        return result
        
    except Exception as e:
        log.error(f"Failed to communicate with Video Agent: {e}")
        raise RuntimeError(f"Video Agent communication failed: {e}")

  
# ---------------- Gradio Main Function (UI) ----------------  
  
def generate_image_gradio(prompt: str, use_local: bool) -> Tuple[Optional[Image.Image], str]:  
    """
    الدالة الرئيسية لواجهة Gradio التفاعلية.  
    """  
    if not prompt:  
        return None, "الرجاء إدخال نص مطالبة (prompt) للتوليد."  
  
    # 1. التوليد المحلي  
    if use_local and LOCAL_DIFFUSERS:  
        return local_generate(prompt)  
      
    # 2. واجهة HF API  
    return hf_inference(prompt)  
  
  
# ---------------- Gradio API Endpoint (For Publisher Agent) ----------------  
  
def gradio_api_endpoint(prompt: str) -> Dict[str, Any]:  
    """
    نقطة النهاية المخصصة للاستدعاء بواسطة Gradio Client (وكيل النشر).  
    """  
    log.info(f"API Endpoint received prompt: {prompt[:50]}...")  
      
    # تفضيل الوضع المحلي على واجهة HF API إذا كان متاحاً  
    if LOCAL_DIFFUSERS:  
        image, status_msg = local_generate(prompt)  
    else:  
        image, status_msg = hf_inference(prompt)  
          
    if image is None:  
        log.error(f"API Endpoint failed to generate image: {status_msg}")  
        raise RuntimeError(f"فشل توليد الصورة: {status_msg}")  
  
    # تحويل صورة PIL إلى Base64  
    buffered = io.BytesIO()  
    image.save(buffered, format="JPEG")  
    img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")  
      
    log.info(f"API Endpoint success. Image size: {len(img_str)} bytes.")  
  
    return {  
        "image_base64": img_str,  
        "status": status_msg  
    }


# ---------------- NEW: Extended API with Video Generation ----------------

def gradio_api_with_video(
    prompt: str,
    quote: str,
    author: str,
    culture: str,
    generate_video: bool = False
) -> Dict[str, Any]:
    """
    نقطة نهاية موسعة تدعم توليد الفيديو اختيارياً.
    
    Args:
        prompt: SD prompt للصورة
        quote: نص الحكمة
        author: قائل الحكمة
        culture: الثقافة
        generate_video: هل يتم توليد فيديو؟
        
    Returns:
        Dict with image_base64, video_base64 (optional), status
    """
    log.info(f"Extended API endpoint called with video={generate_video}")
    
    # 1. توليد الصورة أولاً
    image_result = gradio_api_endpoint(prompt)
    
    response = {
        "image_base64": image_result["image_base64"],
        "image_status": image_result["status"],
        "video_base64": None,
        "video_status": None
    }
    
    # 2. إذا كان مطلوباً توليد فيديو
    if generate_video:
        try:
            video_result = send_to_video_agent(
                image_base64=image_result["image_base64"],
                quote=quote,
                author=author,
                culture=culture
            )
            
            response["video_base64"] = video_result.get("video_base64")
            response["video_status"] = video_result.get("status")
            response["video_path"] = video_result.get("video_path")
            
        except Exception as e:
            log.error(f"Video generation failed: {e}")
            response["video_status"] = f"Video generation failed: {str(e)}"
    
    return response
  
  
# ---------------- Gradio Interface Definition ----------------  
  
with gr.Blocks(title="Image Agent") as demo:  
    gr.Markdown("# 🎨 وكيل توليد الصور (Image Agent)")  
    gr.Markdown(f"**حالة الإعدادات:** {STATUS_MESSAGE}")  
    gr.Markdown("---")  
      
    # ------------------ 1. واجهة المستخدم التفاعلية ------------------  
    with gr.Tab("واجهة المستخدم"):  
        gr.Markdown("## التوليد التفاعلي")  
        with gr.Row():  
            with gr.Column(scale=2):  
                prompt_input = gr.Textbox(  
                    label="نص المطالبة (Prompt) ✍️",  
                    placeholder="كلب آلي فضائي يرتدي بذلة رائد فضاء في عالم مستقبلي، فن رقمي.",  
                    lines=3  
                )  
                  
                # خيار التوليد المحلي (يظهر فقط إذا كان متاحاً)  
                if LOCAL_DIFFUSERS:  
                     local_checkbox = gr.Checkbox(  
                        label="استخدام الوضع المحلي (GPU/CPU)",  
                        value=True,  
                        interactive=True  
                     )  
                else:  
                     local_checkbox = gr.Checkbox(  
                        label="استخدام الوضع المحلي (غير متاح)",  
                        value=False,  
                        interactive=False  
                     )  
                  
                generate_button = gr.Button("🚀 توليد الصورة", variant="primary")  
                  
                status_output = gr.Textbox(label="الحالة والرسائل", max_lines=2)  
  
            with gr.Column(scale=1):  
                image_output = gr.Image(label="الصورة الناتجة", type="pil")  
  
        # ربط الدالة بزر Gradio  
        generate_button.click(  
            fn=generate_image_gradio,  
            inputs=[prompt_input, local_checkbox],  
            outputs=[image_output, status_output]  
        )  
          
    # ------------------ 2. نقطة النهاية للـ API الخارجي ------------------  
    with gr.Tab("نقطة النهاية (API)"):  
        gr.Markdown("## نقطة النهاية لوكيل النشر (API)")  
        gr.Markdown(  
            "هذه الواجهة متاحة للاستدعاء عبر Gradio Client من أنظمة خارجية."  
        )  
          
        # واجهة API بسيطة  
        api_input = gr.Textbox(label="نص المطالبة", placeholder="حكمة فلسفية، فن تجريدي...")  
        api_output = gr.JSON(label="استجابة JSON (image_base64)")  
          
        gr.Button("اختبار API").click(  
            fn=gradio_api_endpoint,  
            inputs=[api_input],  
            outputs=[api_output],  
            api_name="predict"   
        )
    
    # ------------------ 3. نقطة النهاية الموسعة مع الفيديو ------------------
    with gr.Tab("API مع فيديو"):
        gr.Markdown("## 🎬 نقطة نهاية موسعة مع دعم الفيديو")
        gr.Markdown(
            "هذه النقطة تدعم توليد الصورة والفيديو معاً. "
            "مناسبة للاستخدام من وكيل النشر."
        )
        
        with gr.Row():
            with gr.Column():
                ext_prompt = gr.Textbox(label="SD Prompt")
                ext_quote = gr.Textbox(label="Quote")
                ext_author = gr.Textbox(label="Author")
                ext_culture = gr.Textbox(label="Culture")
                ext_video_check = gr.Checkbox(label="Generate Video", value=True)
                
                ext_button = gr.Button("🚀 Generate")
            
            with gr.Column():
                ext_output = gr.JSON(label="Response")
        
        ext_button.click(
            fn=gradio_api_with_video,
            inputs=[ext_prompt, ext_quote, ext_author, ext_culture, ext_video_check],
            outputs=[ext_output],
            api_name="generate_with_video"
        )
  
  
# ---------------- Main Entry Point ----------------  
if __name__ == "__main__":  
    PORT = int(os.getenv("PORT", "7860"))  
    log.info("Starting Image Agent Gradio Interface...")  
    demo.launch(server_name="0.0.0.0", server_port=PORT)