from __future__ import annotations

import asyncio
import glob
import json
import os
import pathlib
import random
import re
import shutil
import subprocess
import sys
import tempfile
import traceback
import uuid
from typing import Any

import base64
import threading
import time

# --- PERSISTENT HARDWARE & RUNTIME ACCELERATION CONFIGS ---
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
os.environ["OMP_NUM_THREADS"] = "4"
os.environ["MKL_NUM_THREADS"] = "4"

import gradio as gr
import requests as http_requests
import spaces
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

if torch.cuda.is_available():
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.backends.cudnn.allow_tf32 = True

# --- SYSTEM DIRECTORIES CONFIGURATION MAPPED TO BUCKET (/data) ---
ROOT = pathlib.Path(__file__).resolve().parent
COMFY = ROOT / "ComfyUI"
INPUT = COMFY / "input"
OUTPUT = COMFY / "output"

# Robust Storage Detection: Use persistent /data volume if available, fallback safely if not.
DATA_DIR = pathlib.Path("/data")
try:
    if not DATA_DIR.exists():
        DATA_DIR.mkdir(parents=True, exist_ok=True)
    # Test file write to confirm permissions are active
    test_dummy = DATA_DIR / f".perm_test_{os.getpid()}"
    test_dummy.touch()
    test_dummy.unlink()
except Exception:
    print("Warning: Persistent storage volume '/data' not accessible. Falling back to local workspace storage.", flush=True)
    DATA_DIR = ROOT / "data"
    DATA_DIR.mkdir(parents=True, exist_ok=True)

MODELS = DATA_DIR / "models"
ENHANCER_DIR = DATA_DIR / "prompt_enhancer"

# Inject the extracted prompt enhancer engine libraries directly into the container's environment paths
os.environ["LD_LIBRARY_PATH"] = f"{ENHANCER_DIR}:{os.environ.get('LD_LIBRARY_PATH', '')}"

WORKFLOW_REPO = "TenStrip/LTX2.3-10Eros_Workflows"
WORKFLOW_REVISION = "1b8e8988842a5850dbba58d732c3e29ce430c1c7"
WORKFLOW_FILENAME = "10Eros_10SNodes_LikenessGuideHelper_I2V_v3.2.json"

# Core Graph Identifiers
RUNEXX_NODE_LOAD_IMAGE_REF2 = 29
RUNEXX_NODE_LOAD_IMAGE_BG = 30
NODE_OUTPUT = "597"
NODE_LOAD_IMAGE = "834"
NODE_POSITIVE = "536"
NODE_NEGATIVE = "537"
NODE_SEED = "524"
NODE_WIDTH = "791"
NODE_HEIGHT = "792"
NODE_LENGTH = "796"

CUSTOM_NODES = [
    ("ComfyUI-GGUF", "https://github.com/city96/ComfyUI-GGUF.git"),
    ("ComfyUI-LTXVideo", "https://github.com/Lightricks/ComfyUI-LTXVideo.git"),
    ("10S-Comfy-nodes", "https://github.com/TenStrip/10S-Comfy-nodes.git"),
    ("ComfyUI-KJNodes", "https://github.com/kijai/ComfyUI-KJNodes.git"),
    ("rgthree-comfy", "https://github.com/rgthree/rgthree-comfy.git"),
    ("ComfyUI-VideoHelperSuite", "https://github.com/Kosinkadink/ComfyUI-VideoHelperSuite.git"),
    ("RES4LYF", "https://github.com/ClownsharkBatwing/RES4LYF.git"),
    ("ComfyUI-Easy-Use", "https://github.com/yolain/ComfyUI-Easy-Use.git"),
    ("ComfyUI-mxToolkit", "https://github.com/Smirnov75/ComfyUI-mxToolkit.git"),
    ("ComfyMath", "https://github.com/evanspearman/ComfyMath.git"),
    ("ComfyUI-Licon-MSR", "https://github.com/liconstudio/ComfyUI-Licon-MSR.git"),
    ("ComfyUI-RMBG", "https://github.com/1038lab/ComfyUI-RMBG.git"),
    ("ComfyUI-PromptRelay", "https://github.com/kijai/ComfyUI-PromptRelay.git"),
    ("ComfyUI-FunPack", "https://github.com/digital-garbage/ComfyUI-FunPack.git"),
    ("ComfyUI-MultiLoRALoader", "https://github.com/phazei/ComfyUI-MultiLoRALoader.git"),
]

# Universal model component configuration mapping directly into persistent directory targets
DOWNLOADS = [
    {"repo": "TenStrip/LTX2.3-10Eros", "file": "10Eros_v1-fp8mixed_learned.safetensors", "dest": MODELS / "checkpoints" / "10Eros_v1-fp8mixed_learned.safetensors"},
    {"repo": "Comfy-Org/ltx-2", "file": "split_files/text_encoders/gemma_3_12B_it_fp8_scaled.safetensors", "dest": MODELS / "text_encoders" / "gemma_3_12B_it_fp8_scaled.safetensors"},
    {"repo": "TenStrip/LTX2.3_Distilled_Lora_1.1_Experiments", "file": "ltx-2.3-22b-distilled-lora-1.1_fro90_ceil72_condsafe.safetensors", "dest": MODELS / "loras" / "ltx23" / "ltx-2.3-22b-distilled-lora-1.1_fro90_ceil72_condsafe.safetensors"},
    {"repo": "VasiliyWeb/OmniNFT_ComfyUI", "file": "OmniNFT_converted_lora.safetensors", "dest": MODELS / "loras" / "ltx23" / "OmniNFT_converted_lora.safetensors"},
    {"repo": "LiconStudio/LTX-2.3-Multiple-Subject-Reference", "file": "LTX2.3-Licon-MSR-test_version.safetensors", "dest": MODELS / "loras" / "ltx23" / "LTX2.3-Licon-MSR-test_version.safetensors"},
    {"repo": "Lightricks/LTX-2.3", "file": "ltx-2.3-spatial-upscaler-x2-1.1.safetensors", "dest": MODELS / "latent_upscale_models" / "ltx-2.3-spatial-upscaler-x2-1.1.safetensors"},
    {"repo": "Kijai/MelBandRoFormer_comfy", "file": "MelBandRoformer_fp16.safetensors", "dest": MODELS / "diffusion_models" / "MelBandRoformer_fp16.safetensors"},
    # Prompt Enhancer Engines saved to permanent /data volume
    {"repo": "signsur4739379373/ltx-dependencies", "file": "llama-server", "dest": ENHANCER_DIR / "llama-server"},
    {"repo": "signsur4739379373/ltx-dependencies", "file": "llama-server-libs.tar.gz", "dest": ENHANCER_DIR / "llama-server-libs.tar.gz"},
    {"repo": "signsur4739379373/ltx-dependencies", "file": "prompt_enhancer/sulphur_prompt_enhancer_model-q8_0.gguf", "dest": ENHANCER_DIR / "sulphur_prompt_enhancer_model-q8_0.gguf"}
]

INPUT_MODES = ["anchor only", "multi-reference (MSR)", "multi-reference (original)"]
PRESETS = ["original", "tuned", "high-fidelity tryon"]
DEFAULT_NEGATIVE = "captions, bad quality, static, low quality, noise, mutant, blur, text, watermark"

_comfy_ready = False
_enhancer_ready = False

def init_services():
    global _comfy_ready, _enhancer_ready
    
    # 1. Base Workspace Cloning
    if not COMFY.exists():
        subprocess.run(["git", "clone", "https://github.com/comfyanonymous/ComfyUI.git", str(COMFY)], check=True)
        
    for name, repo in CUSTOM_NODES:
        node_dir = COMFY / "custom_nodes" / name
        if not node_dir.exists():
            subprocess.run(["git", "clone", repo, str(node_dir)], check=True)

    # 2. Fault-Tolerant Sequential Direct Storage Processing
    for dl in DOWNLOADS:
        dl["dest"].parent.mkdir(parents=True, exist_ok=True)
        if not dl["dest"].exists():
            print(f"Downloading asset to storage bucket: {dl['dest'].name}...", flush=True)
            # FIXED: Added try-except mapping block to survive upstream repository file deletion or 404 errors
            try:
                hf_hub_download(repo_id=dl["repo"], filename=dl["file"], local_dir=dl["dest"].parent)
                # Normalize complex path naming maps if present
                basename = os.path.basename(dl["file"])
                if (dl["dest"].parent / basename).exists() and (dl["dest"].parent / basename) != dl["dest"]:
                    shutil.move(str(dl["dest"].parent / basename), str(dl["dest"]))
            except Exception as download_error:
                print(f"Error skipping non-critical asset download failure for {dl['file']}: {download_error}", flush=True)

    # 3. Dynamic Unpacking and Execution Permissions Modification
    tar_file = ENHANCER_DIR / "llama-server-libs.tar.gz"
    if tar_file.exists() and not (ENHANCER_DIR / "libggml.so").exists():
        print("Extracting prompt enhancer shared systems library binaries...", flush=True)
        try:
            subprocess.run(["tar", "-xzvf", str(tar_file), "-C", str(ENHANCER_DIR)], check=True)
        except Exception as e:
            print(f"Failed parsing enhancer libraries: {e}", flush=True)

    server_bin = ENHANCER_DIR / "llama-server"
    if server_bin.exists():
        os.chmod(str(server_bin), 0o755)

    # 4. Start ComfyUI Server Interface (Optimized for fast swapping)
    if not _comfy_ready:
        print("Starting CPU-Cached ComfyUI Node Server...", flush=True)
        subprocess.Popen([
            sys.executable, str(COMFY / "main.py"), 
            "--listen", "127.0.0.1", "--port", "8188", 
            "--normalvram", "--fp8_e4m3fn-text-enc", "--fp8_e4m3fn-unet"
        ])
        for _ in range(60):
            try:
                if http_requests.get("http://127.0.0.1:8188/object_info").status_code == 200:
                    _comfy_ready = True
                    break
            except Exception:
                time.sleep(1.5)

    # 5. Start Prompt Enhancer Server Process (Only if the model file successfully downloaded)
    model_file = ENHANCER_DIR / "sulphur_prompt_enhancer_model-q8_0.gguf"
    if not _enhancer_ready and server_bin.exists() and model_file.exists():
        print("Starting Prompt Enhancer Text Server Local Engine...", flush=True)
        subprocess.Popen([
            str(server_bin),
            "-m", str(model_file),
            "--port", "8080",
            "--host", "127.0.0.1",
            "-ngl", "0", # Force 0 GPU layers to run efficiently on CPU background thread without consuming core GPU tokens
            "-c", "2048"
        ], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        
        for _ in range(45):
            try:
                if http_requests.get("http://127.0.0.1:8080/health").status_code == 200:
                    _enhancer_ready = True
                    break
            except Exception:
                time.sleep(1.5)

def run_prompt_enhancer(user_prompt: str) -> str:
    init_services()
    if not _enhancer_ready:
        return user_prompt # Return the original prompt cleanly if the enhancer server isn't available
    try:
        payload = {
            "prompt": f"<|im_start|>system\nYou are a precise prompt expansion engine. Expand this raw visual description for high-fidelity rendering details:<|im_end|>\n<|im_start|>user\n{user_prompt}<|im_end|>\n<|im_start|>assistant\n",
            "n_predict": 128,
            "temperature": 0.7
        }
        res = http_requests.post("http://127.0.0.1:8080/completion", json=payload, timeout=10).json()
        return res.get("content", user_prompt).strip()
    except Exception:
        return user_prompt

@spaces.GPU(duration=120)
def process_generation(prompt, neg_prompt, mode, image, ref2, ref3, ref4, background, width, height, length, seed, preset_name):
    init_services()
    if not _comfy_ready: return None, "Backend pipeline initiation failure."

    workflow_path = ROOT / WORKFLOW_FILENAME
    if not workflow_path.exists():
        hf_hub_download(repo_id=WORKFLOW_REPO, filename=WORKFLOW_FILENAME, local_dir=ROOT, revision=WORKFLOW_REVISION)

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

    INPUT.mkdir(parents=True, exist_ok=True)
    OUTPUT.mkdir(parents=True, exist_ok=True)
    run_uuid = str(uuid.uuid4())
    
    if image:
        image.save(INPUT / f"{run_uuid}_ref1.png")
        graph[NODE_LOAD_IMAGE]["widgets_values"][0] = f"{run_uuid}_ref1.png"
    if ref2 and mode != "anchor only":
        ref2.save(INPUT / f"{run_uuid}_ref2.png")
        if str(RUNEXX_NODE_LOAD_IMAGE_REF2) in graph:
            graph[str(RUNEXX_NODE_LOAD_IMAGE_REF2)]["widgets_values"][0] = f"{run_uuid}_ref2.png"
    if background and mode != "anchor only":
        background.save(INPUT / f"{run_uuid}_bg.png")
        if str(RUNEXX_NODE_LOAD_IMAGE_BG) in graph:
            graph[str(RUNEXX_NODE_LOAD_IMAGE_BG)]["widgets_values"][0] = f"{run_uuid}_bg.png"

    w_aligned = (int(width) // 32) * 32
    h_aligned = (int(height) // 32) * 32

    graph[NODE_POSITIVE]["widgets_values"][0] = prompt
    graph[NODE_NEGATIVE]["widgets_values"][0] = neg_prompt if neg_prompt else DEFAULT_NEGATIVE
    graph[NODE_SEED]["widgets_values"][0] = int(seed) if seed != -1 else random.randint(0, 1000000)
    graph[NODE_WIDTH]["widgets_values"][0] = w_aligned
    graph[NODE_HEIGHT]["widgets_values"][0] = h_aligned
    graph[NODE_LENGTH]["widgets_values"][0] = int(length)

    # Fast pruning of audio nodes
    for key in ["274", "535", "550", "617", "789"]:
        if key in graph: del graph[key]

    try:
        res = http_requests.post("http://127.0.0.1:8188/prompt", json={"prompt": graph}).json()
        prompt_id = res["prompt_id"]
    except Exception as e:
        return None, f"Execution dispatch error: {str(e)}"

    while True:
        try:
            history = http_requests.get(f"http://127.0.0.1:8188/history/{prompt_id}").json()
            if prompt_id in history: break
        except Exception:
            pass
        time.sleep(0.5)

    output_files = glob.glob(str(OUTPUT / "*.*"))
    if output_files:
        return max(output_files, key=os.path.getctime), "Generation Complete!"
    
    return None, "Pipeline completed without producing render tracks."

def _on_input_mode_change(m):
    any_msr = m in ["multi-reference (MSR)", "multi-reference (original)"]
    is_msr_ours = m == "multi-reference (MSR)"
    return (
        gr.update(label="reference 1" if any_msr else "reference image"),
        gr.update(visible=any_msr),
        gr.update(visible=is_msr_ours),
        gr.update(visible=is_msr_ours),
        gr.update(visible=any_msr),
    )

with gr.Blocks() as demo:
    gr.Markdown("### ⚡ ZeroGPU Multi-Subject Persistent High-Fidelity Studio")
    
    with gr.Row():
        with gr.Column(scale=4):
            prompt = gr.Textbox(label="Positive Structural Script Prompt", lines=3, placeholder="Describe scene movement...")
            enhance_btn = gr.Button("🔮 Auto-Enhance Description Prompt", variant="secondary")
            neg_prompt = gr.Textbox(label="Negative Restrictions", value=DEFAULT_NEGATIVE, lines=2)
            input_mode = gr.Dropdown(choices=INPUT_MODES, value="anchor only", label="Workflow Input Matrix Configuration Type")
            preset = gr.Dropdown(choices=PRESETS, value="tuned", label="Likeness Precision Profiles")
            
            with gr.Row():
                width = gr.Slider(minimum=256, maximum=1024, value=768, step=32, label="Width")
                height = gr.Slider(minimum=256, maximum=1024, value=512, step=32, label="Height")
                length = gr.Slider(minimum=8, maximum=80, value=41, step=4, label="Frame Length")
                
            seed = gr.Number(label="Manual Seed (-1 = Random)", value=-1, precision=0)
            generate_btn = gr.Button("Render Frame Generation Sequence", variant="primary")
            
        with gr.Column(scale=3):
            image = gr.Image(type="pil", label="reference image")
            msr_ref2 = gr.Image(type="pil", label="reference 2 (Pose/Cloth Asset)", visible=False)
            msr_ref3 = gr.Image(type="pil", label="reference 3", visible=False)
            msr_ref4 = gr.Image(type="pil", label="reference 4", visible=False)
            msr_background = gr.Image(type="pil", label="target background", visible=False)
            
            video_output = gr.Video(label="Processed Result Generation Sequence")
            status_text = gr.Textbox(label="System Status Log", interactive=False)

    input_mode.change(
        fn=_on_input_mode_change,
        inputs=[input_mode],
        outputs=[image, msr_ref2, msr_ref3, msr_ref4, msr_background],
    )

    enhance_btn.click(
        fn=run_prompt_enhancer,
        inputs=[prompt],
        outputs=[prompt]
    )

    generate_btn.click(
        fn=process_generation,
        inputs=[prompt, neg_prompt, input_mode, image, msr_ref2, msr_ref3, msr_ref4, msr_background, width, height, length, seed, preset],
        outputs=[video_output, status_text]
    )

if __name__ == "__main__":
    demo.queue().launch(theme=gr.themes.Soft())