import os
import sys
import requests
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', '')

# تابع ارسال مستقیم خطا به سرور جهت فعال‌سازی سیستم تلاش مجدد (Retry)
def report_failure(error_msg):
    try:
        requests.post(
            f"{space_url}/api/webhook/fail",
            json={
                "run_id": run_id,
                "error": error_msg,
                "event_type": "effects",
                "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 configuration from payload...')
if not raw_prompt.startswith("VOICECONFIG_"):
    err_str = "Error: Invalid configuration payload signature."
    print(err_str)
    report_failure(err_str)
    sys.exit(1)

# استخراج داده‌ها از توکن پیکربندی بدون تداخل با فرآیند ترجمه
config_str = raw_prompt[len("VOICECONFIG_"):]
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)
variant = config.get("variant", "medium")
sampler_type = config.get("sampler", "pingpong")

try:
    duration = float(config.get("duration", "60"))
except ValueError:
    duration = 60.0

try:
    steps = int(config.get("steps", "8"))
except ValueError:
    steps = 8

try:
    cfg_scale = float(config.get("cfg", "1.0"))
except ValueError:
    cfg_scale = 1.0

try:
    seed = int(config.get("seed", "0"))
except ValueError:
    seed = 0

# رمزگشایی پرامپت انگلیسی تولید صدا
try:
    b64_prompt = config.get("prompt", "")
    b64_prompt += "=" * ((4 - len(b64_prompt) % 4) % 4)
    prompt = base64.b64decode(b64_prompt).decode('utf-8')
except Exception as e:
    prompt = ""

print(f"   -> User Run ID: {user_run_id}")
print(f"   -> Model Variant: {variant}")
print(f"   -> Prompt: {prompt}")
print(f"   -> Duration: {duration}s")
print(f"   -> Steps: {steps}")
print(f"   -> CFG Scale: {cfg_scale}")
print(f"   -> Sampler: {sampler_type}")
print(f"   -> Seed: {seed}")

print('2. Connecting to Stable Audio 3 Space...')
try:
    hf_token = os.environ.get('HF_TOKEN', '')
    if hf_token:
        client = Client("stabilityai/stable-audio-3", token=hf_token)
    else:
        client = Client("stabilityai/stable-audio-3")

    print('3. Generating audio from Stable Audio 3...')
    # فراخوانی با آرگومان‌های نام‌گذاری شده برای سازگاری کامل با کلاینت Gradio
    result = client.predict(
        variant_key=variant,
        prompt=prompt,
        duration=duration,
        steps=steps,
        cfg_scale=cfg_scale,
        sampler_type=sampler_type,
        seed=seed,
        api_name="/infer"
    )

    # تحلیل پویای ساختار پاسخ دریافتی جهت استخراج مسیر فایل صوتی نهایی
    audio_path = None
    if isinstance(result, (list, tuple)) and len(result) > 0:
        audio_path = result[0]
    elif isinstance(result, dict):
        audio_path = result.get('name') or result.get('path') or result.get('url')
    else:
        audio_path = result

    # بررسی ساختارهای تو در تو در پاسخ
    if isinstance(audio_path, dict):
        audio_path = audio_path.get('path') or audio_path.get('name') or audio_path.get('url')

    if not audio_path or not os.path.exists(str(audio_path)):
        raise Exception("Generated audio file not found or invalid.")

    print('4. Uploading generated audio back to Server...')
    with open(audio_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': 'wav'}, 
            files={'file': f}
        )

    if res_upload.status_code == 200:
        print('5. 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 audio generation: {err_str}")
    report_failure(err_str)
    sys.exit(1)