import os, sys, random, requests, json, time, base64
import httpx
from PIL import Image

# =========================================================
# رفع قطعی مشکل Time Out رانر گیت‌هاب با تنظیم زمان انتظار به ۵ دقیقه
original_client_init = httpx.Client.__init__
def patched_client_init(self, *args, **kwargs):
    kwargs['timeout'] = httpx.Timeout(300.0)
    original_client_init(self, *args, **kwargs)
httpx.Client.__init__ = patched_client_init

original_async_init = httpx.AsyncClient.__init__
def patched_async_init(self, *args, **kwargs):
    kwargs['timeout'] = httpx.Timeout(300.0)
    original_async_init(self, *args, **kwargs)
httpx.AsyncClient.__init__ = patched_async_init
# =========================================================

from gradio_client import Client

prompt = os.environ.get('PROMPT', '')
image_url = os.environ.get('IMAGE_URL', '')
run_id = os.environ.get('RUN_ID')
space_url = os.environ.get('SPACE_URL')
github_run_id = os.environ.get('GITHUB_RUN_ID')

# دریافت متغیرهای عددی با مقادیر پیش‌فرض هماهنگ با فرانت‌اند جدید
try:
    steps = int(os.environ.get('STEPS', '15'))
except:
    steps = 15

try:
    guidance_scale = float(os.environ.get('GUIDANCE', '1.0'))
except:
    guidance_scale = 1.0

def report_failure(error_msg):
    try:
        requests.post(
            f"{space_url}/api/webhook/fail",
            json={
                "run_id": run_id,
                "error": error_msg,
                "event_type": "generate-editor",
                "client_payload": {
                    "prompt": prompt,
                    "image_url": image_url,
                    "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. Downloading input image from Docker Space...')
local_input_path = 'input_to_edit.png'
try:
    req = requests.get(image_url, timeout=30)
    with open(local_input_path, 'wb') as f:
        f.write(req.content)
    print('   -> Input image ready.')
except Exception as e:
    err_str = f"Download failed: {e}"
    print(err_str)
    report_failure(err_str)
    sys.exit(1)

print('2. Extracting exact original dimensions and preparing Base64 Data URI...')
try:
    # استخراج ابعاد دقیق پیکسلی کاربر برای بازنشانی نهایی
    img = Image.open(local_input_path)
    orig_width, orig_height = img.size
    print(f'   -> Target Exact Size required by User: {orig_width}x{orig_height}')
    
    # تبدیل تصویر دانلود شده به فرمت کد دیتا (Base64 Data URI) همانطور که اسپیس جدید نیاز دارد
    with open(local_input_path, "rb") as img_file:
        encoded_string = base64.b64encode(img_file.read()).decode('utf-8')
    
    base64_data_uri = f"data:image/png;base64,{encoded_string}"
    
    # ساخت ساختار آرایه رشته‌ای جی‌سون (JSON String of Images Array)
    images_json_payload = json.dumps([base64_data_uri])
    print('   -> Base64 payload generated successfully.')

except Exception as e:
    err_str = f"Failed to prepare image data payload: {e}"
    print(err_str)
    report_failure(err_str)
    sys.exit(1)

print('3. Connecting to Qwen Image Edit Space (Fast LoRA)...')
success = False
raw_result = None

for attempt in range(3):
    try:
        print(f'   -> Attempt {attempt + 1} of 3...')
        
        # اتصال به آدرس اسپیس جدید معرفی شده در فرانت‌اند شما
        client = Client("https://prithivmlmods-qwen-image-edit-2511-loras-fast.hf.space/")
        
        seed = random.randint(1, 2147483647)
        lora_adapter = "Anything2Real" # استایل پیش‌فرض برای تبدیل به عکس طبیعی و واقعی
        randomize_seed = True
        
        # ارسال پارامترها دقیقاً منطبق با Predict شماره 2 اسپیس مقصد
        result = client.predict(
            images_json_payload,     # پارامتر اول: رشته جی‌سون تصاویر
            prompt,                  # پارامتر دوم: دستور متنی
            lora_adapter,            # پارامتر سوم: نام آداپتور لورا
            seed,                    # پارامتر چهارم: سید عددی
            randomize_seed,          # پارامتر پنجم: رندوم کردن سید
            guidance_scale,          # پارامتر ششم: مقدار Guidance Scale
            steps,                   # پارامتر هفتم: تعداد مراحل پردازش
            fn_index=2               # مشخص کردن اجرای ایندکس شماره ۲ API
        )
        
        if result:
            raw_result = result
            success = True
            break
    except Exception as client_err:
        print(f'   -> Attempt {attempt + 1} failed: {client_err}')
        time.sleep(5)

if not success or not raw_result:
    err_str = "Qwen-Image-Edit Space failed to process the image after 3 attempts."
    print(f'CRITICAL ERROR: {err_str}')
    report_failure(err_str)
    sys.exit(1)

print('4. Resolving and downloading generated output image...')
local_output_path = 'output_raw.png'
try:
    # استخراج مسیر فایل یا لینک دانلود از پاسخ Gradio Client
    if isinstance(raw_result, (tuple, list)):
        res_data = raw_result[0]
    elif isinstance(raw_result, dict):
        res_data = raw_result.get('data', [None])[0] or raw_result
    else:
        res_data = raw_result

    if isinstance(res_data, dict):
        output_url_or_path = res_data.get('url') or res_data.get('path')
    else:
        output_url_or_path = str(res_data)

    # اگر خروجی یک لینک وب بود آن را دانلود کن، در غیر این صورت فایل محلی را بخوان
    if output_url_or_path.startswith('http'):
        img_req = requests.get(output_url_or_path, timeout=60)
        with open(local_output_path, 'wb') as f:
            f.write(img_req.content)
    else:
        if os.path.exists(output_url_or_path):
            import shutil
            shutil.copy(output_url_or_path, local_output_path)
        else:
            raise FileNotFoundError(f"Result file path not found: {output_url_or_path}")
            
    print('   -> Output downloaded and ready for post-processing.')
except Exception as extract_err:
    err_str = f"Failed to resolve processed image output: {extract_err}. Raw response was: {raw_result}"
    print(err_str)
    report_failure(err_str)
    sys.exit(1)

print('5. Post-Processing: Precisely resizing back to match user original size...')
try:
    # باز کردن تصویر تولید شده جهت اعمال تغییر ابعاد نهایی
    generated_img = Image.open(local_output_path)
    
    # تغییر سایز با بالاترین کیفیت ممکن (فیلتر Lanczos) به ابعاد دقیق پیکسل اولیه کاربر
    final_resized_img = generated_img.resize((orig_width, orig_height), Image.Resampling.LANCZOS)
    
    # ذخیره نهایی بر روی دیسک
    final_resized_img.save(local_input_path) 
    print(f'   -> Done! Output image matches perfectly with original size: {orig_width}x{orig_height} px.')
except Exception as resize_err:
    print(f'   -> Warning: High precision resizing failed ({resize_err}). Using raw output.')
    import shutil
    shutil.copy(local_output_path, local_input_path)

print('6. Preparing image file extension metadata...')
ext = 'png'
try:
    with Image.open(local_input_path) as verify_img:
        content_type = verify_img.format
        if content_type == 'WEBP': ext = 'webp'
        elif content_type in ['JPEG', 'JPG']: ext = 'jpg'
except:
    ext = 'png'

print(f'7. Uploading processed natural photo back to Docker Space: {space_url}')
try:
    with open(local_input_path, 'rb') as f:
        res = requests.post(
            f'{space_url}/api/webhook/upload', 
            data={'run_id': run_id, 'github_run_id': github_run_id, 'ext': ext}, 
            files={'file': f},
            timeout=60
        )
    
    if res.status_code == 200:
        print('8. SUCCESS! Qwen-Image-Edit workflow finished perfectly.')
    else:
        err_str = f"Upload to Docker webhook failed: {res.status_code} - {res.text}"
        print(err_str)
        report_failure(err_str)
        sys.exit(1)
except Exception as up_err:
    err_str = f"Failed to upload result to webhook: {up_err}"
    print(err_str)
    report_failure(err_str)
    sys.exit(1)