import os
import sys
import requests
import json
import base64
from gradio_client import Client

# ۱. دریافت فراداده‌ها و پارامترها از محیط گیت‌هاب رانر اول
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', '')

# مقادیر پیش‌فرض
width = 1024
height = 1024
user_prompt = raw_prompt

print('1. Parsing configuration payload...')
if raw_prompt.startswith("VOICECONFIG_IDEOGRAM_"):
    try:
        # استخراج پرامپت فارسی از انتهای پکیج تنظیمات
        if "_prompt_" in raw_prompt:
            header, b64_prompt = raw_prompt.split("_prompt_", 1)
        else:
            header = raw_prompt
            b64_prompt = ""
            
        config_str = header[len("VOICECONFIG_IDEOGRAM_"):]
        parts = config_str.split("_")
        config = {}
        i = 0
        while i < len(parts) - 1:
            key = parts[i]
            val = parts[i+1]
            if key:
                config[key] = val
            i += 2
            
        width = int(config.get("width", "1024"))
        height = int(config.get("height", "1024"))
        
        # رمزگشایی متن اصلی فارسی از فرمت Base64
        if b64_prompt:
            user_prompt = base64.b64decode(b64_prompt).decode('utf-8')
            print(f"   -> Parsed Persian prompt: {user_prompt}")
    except Exception as parse_err:
        print(f"Error parsing VOICECONFIG: {parse_err}")

def report_failure(error_msg):
    try:
        requests.post(
            f"{space_url}/api/webhook/fail",
            json={
                "run_id": run_id,
                "error": error_msg,
                "event_type": "ideogram",
                "client_payload": {
                    "prompt": raw_prompt,
                    "width": width,
                    "height": height,
                    "run_id": run_id,
                    "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}")

space_name = "ideogram-ai/ideogram4"
print(f"2. Connecting to official Ideogram Space: {space_name}...")

try:
    client = Client(space_name)
    
    # ارسال درخواست ساخت تصویر
    try:
        result = client.predict(
            prompt=user_prompt,
            mode="Quality · 48 steps",
            upsampler="Ideogram (remote)",
            width=width,
            height=height,
            seed=0,
            randomize_seed=True,
            api_name="/generate"
        )
    except Exception:
        result = client.predict(
            user_prompt,
            "Quality · 48 steps",
            "Ideogram (remote)",
            width,
            height,
            0,
            True,
            api_name="/generate"
        )

    # پردازش فایل خروجی
    image_path = None
    raw_json = None

    if isinstance(result, (list, tuple)):
        if len(result) > 0: image_path = result[0]
        if len(result) > 2: raw_json = result[2]
    elif isinstance(result, dict):
        data_list = result.get("data", [])
        if len(data_list) > 0: image_path = data_list[0]
        if len(data_list) > 2: raw_json = data_list[2]
    else:
        image_path = result

    if isinstance(image_path, dict):
        image_path = image_path.get('path') or image_path.get('url')

    if not image_path or not os.path.exists(str(image_path)):
        raise Exception("Output image was not generated or file is missing.")

    print('3. Uploading result raw image back to server...')
    try:
        # آپلود با شناسه _raw جهت تفکیک از فایل نهایی و مخفی ماندن از دید کاربر
        with open(image_path, 'rb') as f:
            res_upload = requests.post(
                f'{space_url}/api/webhook/upload', 
                data={'run_id': f"{run_id}_raw", '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}")

        if raw_json:
            print('4. Uploading JSON metadata details...')
            meta_str = raw_json if isinstance(raw_json, str) else json.dumps(raw_json)
            with open('meta.json', 'w', encoding='utf-8') as mf:
                mf.write(meta_str)
            with open('meta.json', 'rb') as mf_read:
                requests.post(
                    f'{space_url}/api/webhook/upload', 
                    data={'run_id': f'{run_id}_meta', 'github_run_id': github_run_id, 'ext': 'json'}, 
                    files={'file': mf_read},
                    timeout=20
                )

        print('5. SUCCESSFUL!')

    except Exception as up_err:
        raise Exception(f"Upload raw image failed: {up_err}")

except Exception as err:
    err_str = str(err)
    print(f"CRITICAL EXCEPTION OCCURRED: {err_str}")
    report_failure(err_str)
    sys.exit(1)