import os
import uuid
import requests
import threading
import time
from flask import Flask, request, jsonify, render_template, Response
from werkzeug.utils import secure_filename
from itertools import cycle

# --- 1. تنظیمات اصلی ---
WORKER_URLS = [
    "https://ezmary-animat-worker1.hf.space",
    "https://ezmary-animat-worker2.hf.space",
    "https://ezmary-animat-worker3.hf.space",
    "https://ezmary-animat-worker4.hf.space",
    "https://ezmary-animat-worker5.hf.space",
]
worker_pool = cycle(WORKER_URLS)

UPLOADER_API_URL = "https://hamed744-uploadfile.hf.space/upload"
HF_TOKEN = os.getenv("HF_TOKEN")

app = Flask(__name__, template_folder='templates')

jobs = {}
lock = threading.Lock()

def get_next_worker_url():
    return next(worker_pool)

# --- 2. توابع پس‌زمینه ---
def get_permanent_link(job_id, temp_render_url, worker_url):
    try:
        with lock: jobs[job_id]["status"] = "در حال دائمی‌سازی لینک..."
        if not HF_TOKEN: raise Exception("توکن HF برای آپلود دائمی یافت نشد.")
        
        video_full_url = f"{worker_url}{temp_render_url}"
        
        payload = {'url': video_full_url}
        headers = {'Authorization': f'Bearer {HF_TOKEN}'}
        response = requests.post(UPLOADER_API_URL, json=payload, headers=headers, timeout=600)
        response.raise_for_status()
        data = response.json()
        
        uploader_proxy_url = data.get("hf_url")
        if not uploader_proxy_url: raise Exception("پاسخ از سرور آپلودر معتبر نبود.")
        
        path_part = uploader_proxy_url.split('/proxy/')[1]
        final_proxy_url = f"/proxy/{path_part}"
        
        with lock:
            jobs[job_id]["status"] = "completed"
            jobs[job_id]["result"] = final_proxy_url
            
    except Exception as e:
        print(f"خطا در دائمی‌سازی لینک برای {job_id}: {e}")
        with lock:
            jobs[job_id]["status"] = "error"
            jobs[job_id]["result"] = f"خطا در دائمی‌سازی لینک. لینک موقت: {temp_render_url}"

# <<< تابع اصلاح شده: حذف تایم‌اوت و افزایش مقاومت در برابر خطای ارتباطی >>>
def poll_worker_service(job_id, render_job_id, worker_url):
    POLLING_INTERVAL = 15
    MAX_POLLING_ERRORS = 20 # اگر ۲۰ بار پشت سر هم (حدود ۵ دقیقه) نتوانستیم با کارگر صحبت کنیم، خطا می‌دهیم
    error_count = 0

    while True: # حلقه بی‌نهایت، هرگز به خاطر زمان متوقف نمی‌شود
        try:
            response = requests.post(f"{worker_url}/check_job_status", json={"job_id": render_job_id}, timeout=45)
            response.raise_for_status()
            error_count = 0 # ریست کردن شمارنده خطا در صورت موفقیت

            data = response.json()
            with lock:
                current_job = jobs.get(job_id)
                if current_job and current_job.get("status") != data.get("status"):
                    current_job["status"] = data.get("status")
            
            if data.get("status") == "completed":
                get_permanent_link(job_id, data.get("result"), worker_url)
                return # خروج از حلقه
            elif data.get("status") == "error":
                with lock: 
                    current_job = jobs.get(job_id)
                    if current_job:
                        current_job["result"] = data.get("result", "خطای نامشخص از کارگر")
                return # خروج از حلقه
        except requests.exceptions.RequestException as e:
            error_count += 1
            print(f"خطای ارتباطی در پیگیری کارگر {worker_url} (تلاش {error_count}/{MAX_POLLING_ERRORS}): {e}")
            if error_count >= MAX_POLLING_ERRORS:
                with lock:
                    current_job = jobs.get(job_id)
                    if current_job:
                        current_job["status"] = "error"
                        current_job["result"] = f"ارتباط با سرور کارگر ({worker_url}) برای مدت طولانی قطع شد. ممکن است کارگر در حال ری‌استارت باشد. لطفاً بعداً وضعیت را مجدداً بررسی کنید."
                return # خروج از حلقه
        
        time.sleep(POLLING_INTERVAL)

# --- 3. API Endpoints ---
@app.route('/')
def index():
    return render_template('index.html')

@app.route('/api/submit_job', methods=['POST'])
def submit_job():
    if 'image_file' not in request.files or 'video_file' not in request.files:
        return jsonify({"error": "فایل عکس و ویدیو الزامی است."}), 400

    image_file = request.files['image_file']
    video_file = request.files['video_file']
    motion = request.form.get('motion')
    style = request.form.get('style')

    image_bytes = image_file.read()
    video_bytes = video_file.read()
    
    MAX_RETRIES = 3
    for attempt in range(MAX_RETRIES):
        selected_worker_url = get_next_worker_url()
        print(f"تلاش {attempt + 1}/{MAX_RETRIES}: ارسال کار به کارگر: {selected_worker_url}")
        try:
            response = requests.post(
                f"{selected_worker_url}/submit_new_job",
                files={'image_file': (secure_filename(image_file.filename), image_bytes, image_file.mimetype),
                       'video_file': (secure_filename(video_file.filename), video_bytes, video_file.mimetype)},
                data={'motion': motion, 'style': style},
                timeout=600
            )
            response.raise_for_status()
            render_data = response.json()
            render_job_id = render_data.get("job_id")
            if not render_job_id: raise Exception("پاسخ معتبر از کارگر دریافت نشد.")
            
            internal_job_id = str(uuid.uuid4())
            with lock: jobs[internal_job_id] = {"status": "ارسال شده به کارگر...", "result": None}
            
            thread = threading.Thread(target=poll_worker_service, args=(internal_job_id, render_job_id, selected_worker_url))
            thread.start()
            return jsonify({"job_id": internal_job_id})
        except Exception as e:
            print(f"خطا در تلاش {attempt + 1} برای کارگر {selected_worker_url}: {e}")
            time.sleep(2)

    final_error_message = ("<strong>سرورهای پردازش شلوغ هستند!</strong><br>"
                         "لطفاً چند دقیقه دیگر مجدداً تلاش نمایید.")
    return jsonify({"error": final_error_message}), 500

@app.route('/api/check_status', methods=['POST'])
def check_status():
    data = request.get_json()
    job_id = data.get('job_id')
    if not job_id: return jsonify({"error": "شناسه کار الزامی است."}), 400
    with lock: job = jobs.get(job_id, {"status": "not_found", "result": None})
    return jsonify({"status": job["status"], "result": job["result"]})

@app.route('/proxy/<user>/<repo>/<path:file_path>')
def file_proxy(user, repo, file_path):
    if not HF_TOKEN: return "Token سرور برای پراکسی تنظیم نشده است.", 500
    repo_id = f"{user}/{repo}"
    file_url = f"https://huggingface.co/datasets/{repo_id}/resolve/main/{file_path}"
    headers = {"Authorization": f"Bearer {HF_TOKEN}"}
    try:
        hf_response = requests.get(file_url, headers=headers, stream=True, timeout=30)
        hf_response.raise_for_status()
        return Response(hf_response.iter_content(chunk_size=8192),
                        content_type=hf_response.headers.get('Content-Type', 'application/octet-stream'))
    except requests.exceptions.HTTPError as e:
        if e.response.status_code == 404: return "فایل یافت نشد.", 404
        return f"خطای سرور Hugging Face: {e.response.status_code}", e.response.status_code
    except Exception as e:
        print(f"خطای پراکسی: {e}")
        return "خطا در واکشی فایل.", 500

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=int(os.environ.get('PORT', 7860)))