import os
import sys
import requests
import base64
import time
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": "avatar",
                "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("AVATARCONFIG_"):
    err_str = "Error: Invalid avatar configuration payload signature."
    print(err_str)
    report_failure(err_str)
    sys.exit(1)

config_str = raw_prompt[len("AVATARCONFIG_"):]
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")
aud_ext = config.get("audExt", "mp3")
res = config.get("res", "480p")

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

vocal_slug = config.get("vocal", "clean")
vocal_mode = "Clean speech (fast)" if vocal_slug == "clean" else "Isolate vocals (quality)"

accel_slug = config.get("accel", "dbfaster")
if accel_slug == "dbfast":
    acceleration = "DBCache fast"
elif accel_slug == "exact8":
    acceleration = "Exact 8-step"
else:
    acceleration = "DBCache faster"

try:
    b64_prompt = config.get("prompt", "")
    prompt = base64.b64decode(b64_prompt).decode('utf-8')
except Exception as e:
    prompt = "A person is speaking expressively, looking at the camera."

print(f'   -> User Run ID: {user_run_id}')
print(f'   -> Target prompt: "{prompt}"')
print(f'   -> Params: Seed={seed}, Res={res}, Vocal={vocal_mode}, Accel={acceleration}')

# مسیر فایل‌ها در سرور اصلی شما
img_url = f"{space_url}/static/images/{user_run_id}_avatar_img.{img_ext}"
aud_url = f"{space_url}/static/images/{user_run_id}_avatar_aud.{aud_ext}"
state_url = f"{space_url}/static/images/{user_run_id}_state.pt"

local_img = f"input_img.{img_ext}"
local_aud = f"input_aud.{aud_ext}"
local_state = "input_state.pt"

print('2. Downloading assets from your Space...')
# دانلود عکس با تلاش مجدد
for att in range(3):
    try:
        req_img = requests.get(img_url, timeout=45)
        if req_img.status_code == 200:
            with open(local_img, 'wb') as f:
                f.write(req_img.content)
            break
    except:
        time.sleep(2)
else:
    report_failure("Image download failed.")
    sys.exit(1)

# دانلود صدای اصلی با تلاش مجدد
for att in range(3):
    try:
        req_aud = requests.get(aud_url, timeout=45)
        if req_aud.status_code == 200:
            with open(local_aud, 'wb') as f:
                f.write(req_aud.content)
            break
    except:
        time.sleep(2)
else:
    report_failure("Audio download failed.")
    sys.exit(1)

# بررسی وجود فایل پی‌تی (نشان‌دهنده این که این اکشن، ادامه اکشن قبلی است)
has_state = False
try:
    req_state = requests.get(state_url, timeout=45)
    if req_state.status_code == 200:
        with open(local_state, 'wb') as f:
            f.write(req_state.content)
        has_state = True
        print('   -> State file (.pt) found and loaded to resume rendering.')
except Exception as e:
    print("   -> No state file found. Starting fresh generation...")

print('3. Connecting to LongCat-Video-Avatar-1.5 Space on ZeroGPU...')
try:
    client = Client("https://opera8-longcat-video-avatar-1-5-2nd2.hf.space")
    
    # ارتباط با اسپیس هوش مصنوعی با قابلیت تلاش مجدد
    max_submit_retries = 3
    job = None
    for att in range(max_submit_retries):
        try:
            print(f"   -> Submitting job to ZeroGPU (Attempt {att+1})...")
            job = client.submit(
                handle_file(local_img),
                handle_file(local_aud),
                prompt,
                res,
                seed,
                vocal_mode,
                acceleration,
                handle_file(local_state) if has_state else None,
                api_name="/generate"
            )
            break
        except Exception as e:
            print(f"   -> Submission error: {e}")
            time.sleep(5)
    
    if not job:
        raise Exception("Failed to connect to ZeroGPU Space after multiple attempts.")
    
    while not job.done():
        time.sleep(5)
        
    result = job.result()
    video_path = result[0]
    state_path = result[1]
    status_msg = result[2]
    
    # سناریوی اول: ویدیو با موفقیت ساخته شد
    if video_path is not None:
        print('4. Final video generated successfully! Uploading back to your Space...')
        actual_video_path = video_path.get('video') or video_path.get('path') if isinstance(video_path, dict) else video_path
        
        with open(actual_video_path, 'rb') as f:
            res_upload = requests.post(
                f'{space_url}/api/webhook/upload', 
                data={'run_id': user_run_id, 'github_run_id': github_run_id, 'ext': 'mp4'}, 
                files={'file': f},
                timeout=120
            )

        if res_upload.status_code == 200:
            print('5. SUCCESS! Final video process complete.')
            sys.exit(0)
        else:
            raise Exception(f"Webhook upload failed. Status code: {res_upload.status_code}")
            
    # سناریوی دوم: نشست تمام شد و فایل pt دریافت شد (هدایت به اکشن جدید)
    elif state_path is not None:
        print(f"4. Timeout reached securely. State file (.pt) received. Preparing next phase...")
        actual_state_path = state_path.get('path') if isinstance(state_path, dict) else state_path
        
        # آپلود فایل وضعیت جدید برای ذخیره در سرور اصلی
        with open(actual_state_path, 'rb') as f:
            res_state_upload = requests.post(
                f'{space_url}/api/webhook/upload', 
                data={'run_id': f"{user_run_id}_state", 'github_run_id': github_run_id, 'ext': 'pt'}, 
                files={'file': f},
                timeout=120
            )
        
        if res_state_upload.status_code == 200:
            print('   -> State file uploaded. Triggering a NEW GitHub Action runner to continue...')
            
            # ارسال دستور به سرور میانی برای اجرای مجدد یک گیت‌هاب اکشن کاملاً جدید
            dispatch_payload = {
                "prompt": raw_prompt,  # ارسال دقیق همان کانفیگ اولیه (شامل آیدی یکسان)
                "width": 1024,
                "height": 1024,
                "action_name": "avatar"
            }
            
            trigger_res = requests.post(
                f'{space_url}/api/generate',
                json=dispatch_payload,
                timeout=20
            )
            
            if trigger_res.status_code in [200, 204]:
                print('5. SUCCESS! Next phase dispatched to a new Action runner.')
                sys.exit(0)  # خروج کاملا موفق و بدون خطا از اکشن فعلی
            else:
                raise Exception(f"Failed to trigger next phase. Server returned {trigger_res.status_code}")
        else:
            raise Exception(f"Failed to upload state file. Server returned {res_state_upload.status_code}")
    else:
        raise Exception(f"Unexpected result from model: {status_msg}")

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