import os
import sys
import requests
import base64
from gradio_client import Client, handle_file

# دریافت متغیرهای محیطی ارسال شده از سمت گیت‌هاب اکشن
raw_prompt = os.environ.get('PROMPT', '')
run_id = os.environ.get('RUN_ID', '')
space_url = os.environ.get('SPACE_URL', '')
github_run_id = os.environ.get('GITHUB_RUN_ID', '')

# متد گزارش شکست فرآیند جهت فعال‌سازی مجدد رانرهای موازی در صورت بروز هرگونه خطا
def report_failure(error_msg):
    try:
        requests.post(
            f"{space_url}/api/webhook/fail",
            json={
                "run_id": run_id,
                "error": error_msg,
                "event_type": "ghibli",
                "client_payload": {
                    "prompt": raw_prompt,
                    "run_id": run_id,
                    "space_url": space_url
                },
                "github_run_id": github_run_id
            },
            timeout=15
        )
    except Exception as e:
        print(f"Failed to report failure: {e}")

print('1. Decoding packed configurations from safe payload...')
if not raw_prompt.startswith("GHIBLICONFIG_"):
    err_str = "Error: Invalid ghibli configuration payload signature."
    print(err_str)
    report_failure(err_str)
    sys.exit(1)

# استخراج داده‌های رمزگذاری شده کاربر
config_str = raw_prompt[len("GHIBLICONFIG_"):]
parts = config_str.split("_")
config = {}
i = 0
while i < len(parts) - 1:
    key = parts[i]
    val = parts[i+1]
    config[key] = val
    i += 2

user_run_id = config.get("userRunId", run_id)
img_ext = config.get("imgExt", "png")

try:
    b64_style = config.get("style", "")
    style = base64.b64decode(b64_style).decode('utf-8')
except Exception as e:
    style = "Studio Ghibli"

print(f'   -> User Run ID: {user_run_id}')
print(f'   -> Target style: "{style}"')

# آدرس‌دهی و دانلود تصویر پایه از فضای موقت فلاسک
img_url = f"{space_url}/static/images/{user_run_id}_ghibli_img.{img_ext}"
local_img = f"input_img.{img_ext}"

print('2. Downloading reference image from your Space...')
try:
    req_img = requests.get(img_url, timeout=45)
    if req_img.status_code != 200:
        raise Exception(f"Image download failed. Status code: {req_img.status_code}")
    with open(local_img, 'wb') as f:
        f.write(req_img.content)
except Exception as download_err:
    err_str = f"Error downloading source files: {download_err}"
    print(err_str)
    report_failure(err_str)
    sys.exit(1)

# برقراری ارتباط با مدل پشتیبان OminiControl_Art
print('3. Connecting to Yuanshi/OminiControl_Art Space...')
try:
    client = Client("Yuanshi/OminiControl_Art")
    
    print('4. Generating cartoon image style...')
    # ارسال به تابع تولید تصویر با همان پارامترهای دقیق گرادیو
    result = client.predict(
        style,                       # [0] استایل انتخابی کاربر
        handle_file(local_img),      # [1] تصویر ورودی
        "High Quality",              # [2] رزولوشن پردازش
        1.5,                         # [3] غلظت اثرپذیری استایل
        "Auto",                      # [4] پیش‌پردازشگر
        True,                        # [5] رندوم کردن سید خروجی
        42,                          # [6] سید عددی پایه
        20,                          # [7] تعداد گام‌های فرآیند انتشار
        fn_index=7
    )

    # استخراج مسیر فایل ساخته شده در گرادیو کلاینت
    image_path = None
    if isinstance(result, list) and len(result) > 0:
        image_path = result[0].get('image') or result[0].get('path') if isinstance(result[0], dict) else result[0]
    elif isinstance(result, tuple) and len(result) > 0:
        image_path = result[0]
    elif isinstance(result, dict):
        image_path = result.get('image') or result.get('path')
    else:
        image_path = result

    if not image_path or not os.path.exists(str(image_path)):
        raise Exception("Generated output image file not found or invalid.")

    # ارسال تصویر نهایی به شکل وب‌پ فرمت به کنترل‌کننده موقت فلاسک
    print('5. Uploading output image back to your Space...')
    with open(image_path, 'rb') as f:
        res_upload = requests.post(
            f'{space_url}/api/webhook/upload', 
            data={'run_id': run_id, 'github_run_id': github_run_id, 'ext': 'webp'}, 
            files={'file': f}
        )

    if res_upload.status_code == 200:
        print('6. SUCCESS! Process complete.')
    else:
        raise Exception(f"Webhook upload failed. Status code: {res_upload.status_code}")

except Exception as e:
    err_str = str(e)
    print(f"CRITICAL ERROR during generation: {err_str}")
    report_failure(err_str)
    sys.exit(1)