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

# ۱. دریافت فراداده‌ها و پارامترها از گیت‌هاب رانر دوم
raw_prompt = os.environ.get('PROMPT', '')
run_id_for_fail = os.environ.get('RUN_ID', '')
space_url = os.environ.get('SPACE_URL', '')
github_run_id = os.environ.get('GITHUB_RUN_ID', '')

target_run_id = ""
print('1. Parsing upscale configuration payload...')
if raw_prompt.startswith("VOICECONFIG_UPSCALE_"):
    parts = raw_prompt[len("VOICECONFIG_UPSCALE_"):].split("_")
    config = {}
    i = 0
    while i < len(parts) - 1:
        key = parts[i]
        val = parts[i+1]
        if key:
            config[key] = val
        i += 2
    target_run_id = config.get("runId", "")

if not target_run_id:
    target_run_id = run_id_for_fail.replace("_upscale", "")

print(f"   -> Target original run ID to upscale: {target_run_id}")

def report_failure(error_msg):
    try:
        # ارسال شکست برای شناسه جدید تا این بخش ارتقا مجدداً تلاش شود
        requests.post(
            f"{space_url}/api/webhook/fail",
            json={
                "run_id": run_id_for_fail,
                "error": error_msg,
                "event_type": "upscale",
                "client_payload": {
                    "prompt": raw_prompt,
                    "run_id": run_id_for_fail,
                    "space_url": space_url
                },
                "github_run_id": github_run_id
            },
            timeout=10
        )
    except Exception as e:
        print(f"Failed to report failure to main server: {e}")

# دانلود تصویر خام اولیه از سرور
filename_raw = 'input_raw.png'
raw_img_url = f"{space_url}/static/images/{target_run_id}_raw.png"
print(f"2. Downloading raw image from: {raw_img_url}")
try:
    r = requests.get(raw_img_url, timeout=30)
    if r.status_code != 200:
        raise Exception(f"Failed to download raw image. Status: {r.status_code}")
    with open(filename_raw, 'wb') as f:
        f.write(r.content)
except Exception as dl_err:
    err_str = f"Raw image download failed: {dl_err}"
    print(err_str)
    report_failure(err_str)
    sys.exit(1)

# ۳. آغاز فرآیند بهبود و ارتقای کیفیت با استفاده از CodeFormer (فاکتور ۲ و فیدلیتی ۱.۰)
print('3. Connecting to CodeFormer to restore and upscale (Factor=2, Fidelity=1.0)...')
cf_success = False
filename_upscaled = 'output_upscaled.png'

try:
    cf_client = Client('sczhou/CodeFormer')
    upscale_result = None
    
    methods_to_try = [
        {"args": [handle_file(filename_raw), True, True, True, 2, 1.0], "kwargs": {"fn_index": 0}},
        {"args": [handle_file(filename_raw), True, True, 2, 1.0], "kwargs": {"fn_index": 0}},
        {"args": [handle_file(filename_raw), True, True, True, 2, 1.0], "kwargs": {"fn_index": 1}},
        {"args": [handle_file(filename_raw), True, True, 2, 1.0], "kwargs": {"fn_index": 1}},
        {"args": [handle_file(filename_raw), True, True, True, 2, 1.0], "kwargs": {"api_name": "predict"}},
        {"args": [handle_file(filename_raw), True, True, 2, 1.0], "kwargs": {"api_name": "predict"}},
        {"args": [handle_file(filename_raw), True, True, True, 2, 1.0], "kwargs": {"api_name": "/predict"}},
        {"args": [handle_file(filename_raw), True, True, 2, 1.0], "kwargs": {"api_name": "/predict"}},
    ]
    
    for idx, method in enumerate(methods_to_try):
        try:
            print(f'   -> Trying execution path {idx + 1} of {len(methods_to_try)}...')
            upscale_result = cf_client.predict(*method["args"], **method["kwargs"])
            if upscale_result:
                print(f'   -> Success! CodeFormer handled by path {idx + 1}.')
                break
        except Exception as method_err:
            print(f'   -> Path {idx + 1} bypassed: {method_err}')
            
    upscale_path = None
    if isinstance(upscale_result, list) and len(upscale_result) > 0:
        upscale_path = upscale_result[0]
    elif isinstance(upscale_result, tuple) and len(upscale_result) > 0:
        upscale_path = upscale_result[0]
    elif isinstance(upscale_result, dict):
        upscale_path = upscale_result.get('image') or upscale_result.get('path')
    else:
        upscale_path = upscale_result

    if upscale_path and os.path.exists(str(upscale_path)):
        print(f'   -> CodeFormer Success! Output saved at {upscale_path}')
        shutil.copy(upscale_path, filename_upscaled)
        cf_success = True
    else:
        print('   -> Warning: CodeFormer did not return a valid file path.')
except Exception as cf_err:
    print(f'   -> Warning: CodeFormer process failed: {cf_err}')

if not cf_success:
    err_str = "CodeFormer upscaling failed or was bypassed due to connection limits."
    print(f'CRITICAL ERROR: {err_str}')
    report_failure(err_str)
    sys.exit(1)

print('4. Uploading final high-quality image back to server...')
try:
    # آپلود تصویر با شناسه اولیه کاربر تا فرانت‌اند بلافاصله آن را دریافت کند
    with open(filename_upscaled, 'rb') as f:
        res_upload = requests.post(
            f'{space_url}/api/webhook/upload', 
            data={'run_id': target_run_id, 'github_run_id': github_run_id, 'ext': 'png'}, 
            files={'file': f},
            timeout=30
        )
        
    if res_upload.status_code != 200:
        raise Exception(f"Failed to upload image to webhook. Status: {res_upload.status_code}")

    print('5. PROCESS COMPLETED SUCCESSFULLY!')

except Exception as up_err:
    err_str = f"Upload failed: {up_err}"
    print(f"CRITICAL ERROR: {err_str}")
    report_failure(up_err)
    sys.exit(1)