import json, uuid, os, asyncio, time
import httpx
import websockets
from urllib.parse import urlencode
from datetime import datetime

# AMD GPU सर्वर का IP
COMFY_HOST = os.getenv("COMFY_HOST", "134.199.132.159")

with open("workflow.json", "r", encoding="utf-8") as f:
    WORKFLOW_TEMPLATE = json.load(f)

RESOLUTIONS = {
    "360p": {"16:9": (640, 360), "9:16": (360, 640), "1:1": (512, 512), "21:9": (640, 272), "3:4": (384, 512)},
    "480p": {"16:9": (848, 480), "9:16": (480, 848), "1:1": (640, 640), "21:9": (848, 360), "3:4": (480, 640)},
    "720p": {"16:9": (1280, 720), "9:16": (720, 1280), "1:1": (768, 768), "21:9": (1280, 544), "3:4": (768, 1024)}
}

# ==========================================
# 🗄️ SMART HISTORY MANAGER (Live Tracking)
# ==========================================
def load_history():
    if os.path.exists("history.json"):
        try:
            with open("history.json", "r", encoding="utf-8") as f:
                return json.load(f)
        except: pass
    return []

def save_history(data):
    with open("history.json", "w", encoding="utf-8") as f:
        json.dump(data[:50], f, indent=4) # 50 वीडियोस की लिमिट

def init_task(task_id, prompt):
    data = load_history()
    # नया टास्क 'processing' स्टेटस के साथ सबसे ऊपर जोड़ें
    new_entry = {
        "id": task_id,
        "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
        "prompt": prompt,
        "status": "processing",
        "progress": 0,
        "url": None
    }
    data.insert(0, new_entry)
    save_history(data)

def update_task_progress(task_id, progress):
    data = load_history()
    for item in data:
        if item["id"] == task_id:
            item["progress"] = progress
            break
    save_history(data)

def complete_task(task_id, video_url=None, error=None):
    data = load_history()
    for item in data:
        if item["id"] == task_id:
            if error:
                item["status"] = "failed"
                item["error"] = error
            else:
                item["status"] = "done"
                item["progress"] = 100
                item["url"] = video_url
            break
    save_history(data)

# ==========================================
# ⚙️ CORE ENGINE LOGIC
# ==========================================
def inject_params(req: dict) -> dict:
    p = json.loads(json.dumps(WORKFLOW_TEMPLATE))
    p["89"]["inputs"]["text"] = req.get("prompt", "")
    ratio = req.get("aspect_ratio", "16:9")
    quality = req.get("quality", "480p")
    width, height = RESOLUTIONS.get(quality, {}).get(ratio, (848, 480))
    p["74"]["inputs"]["width"] = width
    p["74"]["inputs"]["height"] = height
    p["74"]["inputs"]["length"] = req.get("frames", 81)
    p["88"]["inputs"]["fps"] = req.get("fps", 16)
    
    # ❌ यहाँ से मैंने steps और cfg वाली दोनों लाइनें डिलीट कर दी हैं!
    # अब यह कोड सीधा workflow.json की मास्टर सेटिंग्स उठाएगा।
    return p

def extract_video_url(history: dict, token: str) -> str:
    outputs = history.get("outputs", {})
    for _, node_out in outputs.items():
        for key in ("videos", "images", "files"):
            if key in node_out and node_out[key]:
                it = node_out[key][0]
                q = urlencode(it)
                return f"/api/video?{q}&token={token}"
    raise RuntimeError("Video not found")

async def queue_prompt(req: dict):
    token = str(uuid.uuid4())
    client_id = str(uuid.uuid4())
    prompt_data = inject_params(req)
    
    async with httpx.AsyncClient() as client:
        resp = await client.post(f"http://{COMFY_HOST}/prompt?token={token}", 
                               json={"prompt": prompt_data, "client_id": client_id}, 
                               timeout=30.0)
        prompt_id = resp.json().get("prompt_id")
        if not prompt_id: raise Exception("Failed to queue on AMD Server")
        
    return prompt_id, client_id, token, len(prompt_data)

# 🚀 THE BACKGROUND WATCHER (यह फोन बंद होने पर भी चलता रहेगा)
async def background_watcher(task_id, client_id, token, total_nodes):
    seen = set()
    start_t = time.time()
    progress_fake = 0
    ws_url = f"ws://{COMFY_HOST}/ws?clientId={client_id}&token={token}"
    
    try:
        async with websockets.connect(ws_url, ping_interval=20) as ws:
            while True:
                msg_raw = await ws.recv()
                if isinstance(msg_raw, (bytes, bytearray)):
                    if progress_fake < 95 and (time.time() - start_t) > 2:
                        progress_fake = min(95, progress_fake + 1)
                        update_task_progress(task_id, progress_fake) # लाइव प्रोग्रेस सेव
                    continue
                
                msg = json.loads(msg_raw)
                if msg.get("type") == "executing":
                    node = msg.get("data", {}).get("node")
                    if node is None: break # रेंडरिंग ख़तम
                    if node not in seen:
                        seen.add(node)
                        p_real = int((len(seen) / total_nodes) * 100)
                        progress_fake = max(progress_fake, p_real)
                        update_task_progress(task_id, progress_fake) # लाइव प्रोग्रेस सेव
    except Exception:
        pass # अगर कनेक्शन टूटा तो भी कोई बात नहीं, हम हिस्ट्री चेक कर लेंगे!

    # 🎬 फाइनल चेक और सेविंग
    await asyncio.sleep(2)
    try:
        async with httpx.AsyncClient() as client:
            h_resp = await client.get(f"http://{COMFY_HOST}/history/{task_id}?token={token}", timeout=30.0)
            history = h_resp.json().get(task_id, {})
            v_url = extract_video_url(history, token)
            complete_task(task_id, video_url=v_url) # सक्सेस!
    except Exception as e:
        complete_task(task_id, error=str(e)) # फ़ैल हुआ तो एरर सेव करो