import json
import os
import time
import uuid
import tempfile
from PIL import Image, ImageDraw, ImageFont
import gradio as gr
import base64
import mimetypes
from google import genai
from google.genai import types
import concurrent.futures
import random
from typing import List, Tuple, Optional
import threading

def save_binary_file(file_name, data):
    with open(file_name, "wb") as f:
        f.write(data)

def translate_prompt_to_english(text, api_key, model="gemini-2.0-flash-exp"):
    is_english = all(ord(char) < 128 for char in text)
    if is_english:
        return text

    client = genai.Client(api_key=api_key.strip())

    pre_prompt = "Translate this to English:"
    full_text = pre_prompt + "\n" + text

    contents = [
        types.Content(
            role="user",
            parts=[
                types.Part.from_text(text=full_text),
            ],
        ),
    ]
    generate_content_config = types.GenerateContentConfig(
        temperature=0.5,
        top_p=0.9,
        top_k=40,
        max_output_tokens=8192,
        response_mime_type="text/plain",
    )

    text_response = ""
    for chunk in client.models.generate_content_stream(
        model=model,
        contents=contents,
        config=generate_content_config,
    ):
        if not chunk.candidates or not chunk.candidates[0].content or not chunk.candidates[0].content.parts:
            continue
        text_response += chunk.text + "\n"

    return text_response.strip()

def generate_with_api(api_key, text, file_name, model="gemini-2.0-flash-exp"):
    client = genai.Client(api_key=api_key.strip())
    files = [client.files.upload(file=file_name)]

    pre_prompt = "Apply these changes to the image:"
    full_text = pre_prompt + "\n" + text

    contents = [
        types.Content(
            role="user",
            parts=[
                types.Part.from_uri(
                    file_uri=files[0].uri,
                    mime_type=files[0].mime_type,
                ),
                types.Part.from_text(text=full_text),
            ],
        ),
    ]
    generate_content_config = types.GenerateContentConfig(
        temperature=1,
        top_p=0.95,
        top_k=40,
        max_output_tokens=8192,
        response_modalities=["image", "text"],
        response_mime_type="text/plain",
    )

    text_response = ""
    image_path = None
    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
        temp_path = tmp.name
        for chunk in client.models.generate_content_stream(
            model=model,
            contents=contents,
            config=generate_content_config,
        ):
            if not chunk.candidates or not chunk.candidates[0].content or not chunk.candidates[0].content.parts:
                continue
            candidate = chunk.candidates[0].content.parts[0]
            if candidate.inline_data:
                save_binary_file(temp_path, candidate.inline_data.data)
                image_path = temp_path
                break
            else:
                text_response += chunk.text + "\n"

    del files
    return image_path, text_response

def process_single_api(api_key, prompt, file_name, model):
    if not api_key:
        return None, "API key not provided"
    
    try:
        translated_prompt = translate_prompt_to_english(prompt, api_key, model)
        image_path, text_response = generate_with_api(api_key, translated_prompt, file_name, model)
        
        if image_path:
            result_img = Image.open(image_path)
            if result_img.mode == "RGBA":
                result_img = result_img.convert("RGB")
            return result_img, ""
        return None, text_response if text_response else "No image generated"
    
    except Exception as e:
        return None, f"Error with API {api_key[-4:]}: {str(e)}"

def process_image_and_prompt(composite_pil, prompt):
    try:
        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
            composite_path = tmp.name
            composite_pil.save(composite_path)

        model = "gemini-2.0-flash-exp"
        api_keys = []
        for i in range(1, 51):  # از API_KEY_1 تا API_KEY_50
            key = os.environ.get(f"GEMINI_API_KEY_{i}")
            if key and key.strip():
                api_keys.append(key.strip())
        
        if not api_keys:
            os.unlink(composite_path)
            return None, "هیچ API Key معتبری یافت نشد. لطفاً تنظیمات را بررسی کنید."
        
        # انتخاب تصادفی 4 API از لیست
        random.shuffle(api_keys)
        selected_apis = api_keys[:4]
        
        result_images = []
        error_messages = []
        
        with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor:
            futures = {
                executor.submit(
                    process_single_api,
                    api_key, prompt, composite_path, model
                ): api_key for api_key in selected_apis
            }
            
            for future in concurrent.futures.as_completed(futures):
                image, error = future.result()
                if image:
                    result_images.append(image)
                if error:
                    error_messages.append(error)

        os.unlink(composite_path)

        if not result_images:
            error_text = "\n".join(error_messages) + "\n\n**توجه**: اگر تصویر تولید نشد، لطفاً دستور خود را واضح‌تر بنویسید یا دوباره امتحان کنید."
            return None, error_text
        
        return result_images, ""

    except Exception as e:
        return None, f"خطا در پردازش: {e}"

# تابع جدید برای ریست خودکار هر 24 ساعت
def auto_restart():
    while True:
        time.sleep(24 * 60 * 60)  # 24 ساعت
        print("ریست خودکار Space پس از 24 ساعت")
        os._exit(1)

# شروع ترد ریست خودکار
restart_thread = threading.Thread(target=auto_restart, daemon=True)
restart_thread.start()

css = """
footer { visibility: hidden; }
[class="flagging"], [id*="flagging"],
[class*="gradio-flagging"], [class*="svelte"][class*="flagging"] {
display: none !important;
}
"""

with gr.Blocks(css_paths="style.css", css=css) as demo:
    gr.HTML(
    """
    <div class="header-container">
    <div>
    <img src="https://uploadkon.ir/uploads/4a3e22_25IMG-%DB%B2%DB%B0%DB%B2%DB%B5%DB%B0%DB%B3%DB%B2%DB%B2-%DB%B1%DB%B7%DB%B1%DB%B8%DB%B5%DB%B2.jpg" alt="Alfa AI logo">
    </div>
    <div>
    <h1>ویرایش جادویی تصاویر✨ با هوش مصنوعی آلفا</h1>
    </div>
    </div>
    """
    )

    with gr.Accordion("⚠️ راهنمای استفاده", open=False, elem_classes="config-accordion"):
        gr.Markdown("""
        ### راهنمای استفاده
        - تصویر خود را آپلود کرده و دستور ویرایش را وارد کنید
        - در صورت بروز خطا، پیام مربوطه نمایش داده خواهد شد
        - فقط تصاویر با فرمت PNG آپلود کنید
        - از آپلود تصاویر نامناسب خودداری کنید
        """)

    with gr.Accordion("📌 دستورالعمل‌های ویرایش", open=False, elem_classes="instructions-accordion"):
        gr.Markdown("""
        ### نمونه دستورات ویرایش
        - متن تصویر را به \"متن جدید\" تغییر بده
        - شیء خاصی را از تصویر حذف کن
        - استایل خاصی به بخشی از تصویر اضافه کن
        - تغییرات رنگی روی تصویر اعمال کن
        """)

    with gr.Row(elem_classes="main-content"):
        with gr.Column(elem_classes="input-column"):
            image_input = gr.Image(
                type="pil",
                label="تصویر را آپلود کنید",
                image_mode="RGBA",
                elem_id="image-input",
                elem_classes="upload-box"
            )
            prompt_input = gr.Textbox(
                lines=2,
                placeholder="تصویر چیکار بشه؟ اینجا بنویسید...",
                label="دستور ویرایش",
                elem_classes="prompt-input"
            )
            submit_btn = gr.Button("اعمال تغییرات", elem_classes="generate-btn")
        
        with gr.Column(elem_classes="output-column"):
            output_gallery = gr.Gallery(label="تصاویر ویرایش شده", elem_classes="output-gallery")
            output_text = gr.Textbox(
                label="پیام سیستم", 
                placeholder="در صورت بروز خطا، پیام مربوطه اینجا نمایش داده می‌شود.",
                elem_classes="output-text"
            )

    submit_btn.click(
        fn=process_image_and_prompt,
        inputs=[image_input, prompt_input],
        outputs=[output_gallery, output_text],
    )

    gr.Markdown("## نمونه‌های آماده", elem_classes="gr-examples-header")

    examples = [
        ["data/Negar_1745234748135.jpg", 'متن را به "آلفا" تغییر بده', ""],
        ["data/Negar_1745235002412.jpg", "قاشق را از دست حذف کن", ""],
        ["data/Negar_1745235101190.jpg", 'متن را به "بساز" تغییر بده', ""],
        ["data/1.jpg", "فقط روی صورت استایل جوکر اضافه کن", ""],
        ["data/1777043.jpg", "فقط روی صورت استایل جوکر اضافه کن", ""],
        ["data/2807615.jpg", "فقط روی لب‌ها رژ لب اضافه کن", ""],
        ["data/76860.jpg", "فقط روی لب‌ها رژ لب اضافه کن", ""],
        ["data/2807615.jpg", "فقط صورت را شادتر کن", ""],
    ]

    gr.Examples(
        examples=examples,
        inputs=[image_input, prompt_input],
        elem_id="examples-grid"
    )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)