import os
import sys
import json
import base64
import requests
import random
import time
import io

# ثبت زمان دقیق شروع اجرای اسکریپت رانر برای پیشگیری از رفتارهای پیش‌بینی نشده
start_time = time.time()

# =========================================================
# تنظیم زمان انتظار ارتباطی کلاینت به ۲۵۰ ثانیه جهت پایش مداوم استریم‌های ناپایدار
import httpx
original_client_init = httpx.Client.__init__
def patched_client_init(self, *args, **kwargs):
    kwargs['timeout'] = httpx.Timeout(250.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(250.0)
    original_async_init(self, *args, **kwargs)
httpx.AsyncClient.__init__ = patched_async_init
# =========================================================

from gradio_client import Client, handle_file
from PIL import Image

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', '')

# بررسی و رمزگشایی متغیرهای ارسال شده
if not raw_prompt.startswith("VOICECONFIG_LTX_"):
    err_str = "Error: Invalid LTX configuration payload."
    print(err_str)
    sys.exit(1)

try:
    encoded_data = raw_prompt[len("VOICECONFIG_LTX_"):]
    config = json.loads(base64.b64decode(encoded_data).decode('utf-8'))
    # ذخیره و تثبیت آیدی اولیه کاربر برای پایش مرورگر
    original_run_id = config.get("original_run_id", run_id)
except Exception as e:
    err_str = f"Decode Error: {e}"
    print(err_str)
    sys.exit(1)

def report_failure(error_msg):
    for webhook_attempt in range(3):
        try:
            res = requests.post(
                f"{space_url}/api/webhook/fail",
                json={
                    "run_id": run_id,
                    "error": error_msg,
                    "event_type": "ltxvideo",
                    "client_payload": {
                        "prompt": raw_prompt,
                        "run_id": run_id,
                        "space_url": space_url
                    },
                    "github_run_id": github_run_id
                },
                timeout=20
            )
            print(f"   -> Webhook response status: {res.status_code} - Server Response: {res.text}", flush=True)
            if res.status_code == 200:
                # در صورتی که سرور اصلی اعلام کرد فرآیند ۲۰ تایی به سقف خود رسیده است
                if "Max retries reached" in res.text:
                    try:
                        cycle_count = int(config.get("cycle_count", 1))
                        
                        # شروع دور جدید تلاش (تا ۳ دور، مجموعاً ۶۰ تلاش)
                        if cycle_count < 3:
                            next_cycle = cycle_count + 1
                            print(f"🔄 Cycle {cycle_count} failed after 20 retries. Automatically starting Cycle {next_cycle} out of 3...", flush=True)
                            
                            # به‌روزرسانی شمارنده دور تلاش و حفظ آیدی اولیه کاربر در بدنه پیکربندی
                            config["cycle_count"] = next_cycle
                            config["original_run_id"] = original_run_id
                            
                            new_encoded = base64.b64encode(json.dumps(config).encode('utf-8')).decode('utf-8')
                            new_raw_prompt = f"VOICECONFIG_LTX_{new_encoded}"
                            
                            # آپلود فایل نشانگر وضعیت پردازش پس‌زمینه برای مرورگر کاربر
                            with open("processing.txt", "w", encoding="utf-8") as pf:
                                pf.write("processing")
                            requests.post(
                                f"{space_url}/api/webhook/upload",
                                data={'run_id': f"{original_run_id}_processing", 'github_run_id': github_run_id, 'ext': 'txt'},
                                files={'file': open("processing.txt", "rb")},
                                timeout=30
                            )
                            
                            spawn_res = requests.post(
                                f"{space_url}/api/generate",
                                json={
                                    "prompt": new_raw_prompt,
                                    "action_name": "ltxvideo"
                                },
                                timeout=30
                            )
                            print(f"   -> Spawned new cycle status: {spawn_res.status_code} - Server Resp: {spawn_res.text}", flush=True)
                        else:
                            # اگر پس از ۶۰ تلاش کل فرآیند شکست خورد، وضعیت را خاتمه یافته اعلام کرده و کارت خطا را مپ می‌کنیم
                            print(f"❌ Max cycles reached ({cycle_count}/3). Mapping final error card back to original ID: {original_run_id}", flush=True)
                            
                            with open("processing.txt", "w", encoding="utf-8") as pf:
                                pf.write("failed")
                            requests.post(
                                f"{space_url}/api/webhook/upload",
                                data={'run_id': f"{original_run_id}_processing", 'github_run_id': github_run_id, 'ext': 'txt'},
                                files={'file': open("processing.txt", "rb")},
                                timeout=30
                            )
                            
                            time.sleep(5)  # زمان کوتاه برای اطمینان از ذخیره کامل فایل خطا روی سرور
                            error_card_url = f"{space_url}/static/images/{run_id}.png"
                            err_req = requests.get(error_card_url, timeout=20)
                            if err_req.status_code == 200:
                                requests.post(
                                    f"{space_url}/api/webhook/upload",
                                    data={'run_id': original_run_id, 'github_run_id': github_run_id, 'ext': 'png'},
                                    files={'file': ('error.png', err_req.content)},
                                    timeout=30
                                )
                                print(f"   -> Successfully mapped error card to {original_run_id}.png", flush=True)
                    except Exception as spawn_err:
                        print(f"⚠️ Failed to spawn next cycle automatically: {spawn_err}", flush=True)
                break
        except Exception as e:
            print(f"   -> Webhook failed to send, retrying... {e}", flush=True)
            time.sleep(3)

print(f"🚀 Starting LTX Video runner. Run ID: {run_id} (Original User ID: {original_run_id})")

try:
    client = Client("Fighterdan/LTX-2.3-10Eros_I2V")
except Exception as e:
    err_str = f"Failed to connect to Gradio Space: {e}"
    print(err_str)
    report_failure(err_str)
    sys.exit(1)

# مدیریت دانلود تصویر با ۵ بار تلاش جهت پیشگیری از خطای تاخیر دیسپچ داکر
image_url = config.get("image_path")
local_img = "gen_input.png"
download_success = False
max_download_attempts = 5

print("1. Downloading input image...")
for dl_attempt in range(max_download_attempts):
    try:
        r = requests.get(image_url, timeout=45)
        if r.status_code == 200:
            img_data = io.BytesIO(r.content)
            img = Image.open(img_data)
            img.verify()
            
            img_data.seek(0)
            clean_img = Image.open(img_data).convert('RGB')
            clean_img.save(local_img, format='PNG')
            download_success = True
            print("   -> Image successfully downloaded and verified.")
            break
        else:
            print(f"   -> Download attempt {dl_attempt + 1} failed with status: {r.status_code}")
    except Exception as e:
        print(f"   -> Download attempt {dl_attempt + 1} connection failed: {e}")
    
    if dl_attempt < max_download_attempts - 1:
        print("   -> Waiting 8 seconds before retrying image download...")
        time.sleep(8.0)

if not download_success:
    err_str = "ZeroGPU quota hit during image download. Aborting and switching runner."
    print(f"CRITICAL ERROR: {err_str}", flush=True)
    report_failure(err_str)
    time.sleep(3)
    sys.exit(1)

base_seed = int(config.get("seed", -1))
seconds = float(config.get("seconds", 5.0))
prompt = config.get("prompt", "")

print(f"🎬 Generation configurations - Duration: {seconds}s")

video_path = None
max_attempts = 4

# حلقه تلاش مجدد داخلی برای ساخت ویدیو
for attempt in range(max_attempts):
    elapsed_time = time.time() - start_time
    if elapsed_time > 450.0:
        print(f"⏳ WARNING: Elapsed time is {elapsed_time:.1f}s. Aborting this run to trigger a clean retry on a new runner.")
        report_failure("TIMEOUT_VOLUNTARY_ABORT")
        sys.exit(1)

    seed = base_seed
    if seed == -1:
        seed = random.randint(1, 2147483647)

    print(f"   -> Generation Attempt {attempt + 1} of {max_attempts} with Seed: {seed}...")
    
    try:
        result = client.predict(
            image_path=handle_file(local_img),
            prompt=prompt,
            negative_prompt="captions, music, transition, bad quality, subtitles, text, watermark, cartoon, ugly, blur, static, noise, mutant, horror",
            seconds=seconds,
            preset="tuned",
            seed=seed,
            randomize_seed=False,
            max_width=1120,
            max_height=1344,
            target_mp=1.15,
            snap_multiple=64,
            custom_res_enabled=False,
            mode="anchor only",
            face_bbox="",
            likeness_strength=0.9,
            likeness_anchor_strength=0.15,
            latent_anchor_strength=0.08,
            first_frame_strength=0.82,
            sulphur_lora_strength=0.15,
            sulphur_v1_lora_strength=0.15,
            vbvr_lora_strength=0.5,
            dreamly_lora_strength=0.6,
            synth_lora_strength=0.0,
            plora_lora_strength=0.0,
            singularity_lora_strength=0.3,
            omninft_lora_strength=0.8,
            omninft_bf16_lora_strength=0.0,
            better_motion_lora_strength=0.0,
            physics_v2_lora_strength=0.0,
            hardcut_lora_strength=0.0,
            transition_lora_strength=0.15,
            sulphur_audio_strength=0.15,
            sulphur_v1_audio_strength=0.15,
            vbvr_audio_strength=0.5,
            dreamly_audio_strength=0.6,
            synth_audio_strength=0.0,
            plora_audio_strength=0.0,
            singularity_audio_strength=0.3,
            omninft_audio_strength=0.8,
            omninft_bf16_audio_strength=0.0,
            better_motion_audio_strength=0.0,
            physics_v2_audio_strength=0.0,
            hardcut_audio_strength=0.0,
            transition_audio_strength=0.0,
            cache_at_step=0,
            cache_warmup=400,
            energy_threshold=0.3,
            anchor_similarity_threshold=0.3,
            sigma_string="0.4824, 0.2412, 0.0",
            skip_refine=False,
            gen_budget=0,
            profile_name="",
            input_mode="single image (i2v)",
            msr_ref2=None,
            msr_ref3=None,
            msr_ref4=None,
            msr_background=None,
            msr_frame_count=41,
            msr_guide_strength=1.0,
            msr_lora_strength=0.7,
            prompt_relay_enabled=False,
            prompt_segments="",
            scene_chain_enabled=False,
            scene_chain_prompt="",
            scene_chain_max_scenes=2,
            scene_chain_frame_overlap=8,
            scene_chain_mid_guide=True,
            scene_chain_mid_guide_strength=0.25,
            kv_enabled=False,
            kv_strength=1.0,
            audio_ref_enabled=False,
            audio_ref_file=None,
            audio_ref_guidance_scale=3.0,
            audio_ref_stem_sep=False,
            audio_ref_normalize=True,
            kf_strength=0.82,
            kf_last_image=None,
            kf_mid_enabled=False,
            kf_mid_1_image=None,
            kf_mid_1_pos=50,
            kf_mid_2_image=None,
            kf_mid_2_pos=50,
            kf_mid_3_image=None,
            kf_mid_3_pos=50,
            kf_mid_4_image=None,
            kf_mid_4_pos=50,
            kf_mid_5_image=None,
            kf_mid_5_pos=50,
            api_name="/generate"
        )

        if isinstance(result, (tuple, list)):
            video_data = result[0]
        else:
            video_data = result

        if isinstance(video_data, dict):
            video_path = video_data.get('video') or video_data.get('path')
        else:
            video_path = str(video_data)

        if video_path and os.path.exists(video_path):
            print("   -> Success! Video file created successfully.")
            break

    except Exception as e:
        err_msg = str(e)
        print(f"   -> Failed on attempt {attempt + 1}: {err_msg}")
        
        if 'ZeroGPU quota' in err_msg or 'quota' in err_msg.lower():
            print("   -> ZeroGPU quota hit. Aborting internal attempts to trigger webhook fail and switch runner.")
            report_failure(err_msg)
            sys.exit(1)
            
        if attempt < max_attempts - 1:
            print("   -> Waiting 5 seconds before next internal attempt...")
            time.sleep(5.0)

if not video_path or not os.path.exists(video_path):
    err_str = "Failed to generate video after all internal attempts."
    print(f"CRITICAL ERROR: {err_str}")
    report_failure(err_str)
    sys.exit(1)

# آپلود فایل نهایی تحت آیدی اولیه کاربر جهت ردیابی مستقیم فرانت‌اند
print(f"Uploading video back to space using ID: {original_run_id}...")
ext = 'mp4'
if video_path.endswith('.webm'): 
    ext = 'webm'

try:
    # تغییر وضعیت فایل موقت پایش فرانت‌اند به اتمام پردازش موفقیت‌آمیز
    with open("processing.txt", "w", encoding="utf-8") as pf:
        pf.write("done")
    requests.post(
        f"{space_url}/api/webhook/upload",
        data={'run_id': f"{original_run_id}_processing", 'github_run_id': github_run_id, 'ext': 'txt'},
        files={'file': open("processing.txt", "rb")},
        timeout=30
    )

    with open(video_path, 'rb') as f:
        requests.post(
            f"{space_url}/api/webhook/upload", 
            data={'run_id': original_run_id, 'github_run_id': github_run_id, 'ext': ext}, 
            files={'file': f},
            timeout=60
        )

    print("Uploading Seed Metadata...")
    meta_data = {"seed": seed}
    with open("meta.json", "w", encoding="utf-8") as mf:
        json.dump(meta_data, mf)
        
    with open("meta.json", "rb") as mf:
        requests.post(
            f"{space_url}/api/webhook/upload",
            data={'run_id': f"{original_run_id}_meta", 'github_run_id': github_run_id, 'ext': 'json'},
            files={'file': mf},
            timeout=30
        )

    print("✅ SUCCESS!")

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