import os
import sys
import subprocess
import time
import json
import urllib.request
import urllib.parse
import gradio as gr
import torch
import spaces

# ==========================================
# 1. Install Dependencies Natively at Startup
# ==========================================
required_packages = [
    "safetensors", "scipy", "tqdm", "psutil", "einops", 
    "transformers", "tokenizers", "sentencepiece", "torchsde", 
    "huggingface-hub", "aiohttp", "yarl", "av", "blake3",
    "sqlalchemy", "alembic", "comfy-aimdo"
]

print("Checking system requirements...")
for package in required_packages:
    try:
        __import__(package.replace("-", "_"))
    except ImportError:
        print(f"Installing missing dependency: {package}")
        subprocess.run([sys.executable, "-m", "pip", "install", package], check=True)

# ==========================================
# 2. Clone ComfyUI & Download Custom Model
# ==========================================
COMFYUI_DIR = os.path.abspath("ComfyUI")
if not os.path.exists(COMFYUI_DIR):
    print("Cloning ComfyUI framework...")
    subprocess.run(["git", "clone", "https://github.com/comfyanonymous/ComfyUI.git", COMFYUI_DIR], check=True)

# Ensure the models directory exists
CHECKPOINT_DIR = os.path.join(COMFYUI_DIR, "models", "checkpoints")
os.makedirs(CHECKPOINT_DIR, exist_ok=True)

# Download the exact model your workflow requires
model_filename = "epicphotogasm_ultimateFidelity.safetensors"
model_path = os.path.join(CHECKPOINT_DIR, model_filename)

if not os.path.exists(model_path):
    print(f"Downloading {model_filename} (This may take a few minutes)...")
    # Public HuggingFace mirror for the EpicPhotogasm checkpoint
    model_url = "https://huggingface.co/sibylexpe/ModelsSD15/resolve/main/epicphotogasm_ultimateFidelity.safetensors"
    subprocess.run(["wget", "-q", "-O", model_path, model_url], check=True)


# ==========================================
# 3. Dynamic GPU Inference Function
# ==========================================
# FIX: Capitalized GPU here
@spaces.GPU(duration=120)
def generate_image(user_prompt):
    if not os.path.exists("workflow_api.json"):
        print("Error: workflow_api.json missing from root directory.")
        return None  
        
    with open("workflow_api.json", "r") as f:
        prompt_workflow = json.load(f)

    # Automatically patch the SD3 vs SD 1.5 Latent Mismatch
    if "68" in prompt_workflow and prompt_workflow["68"].get("class_type") == "EmptySD3LatentImage":
        prompt_workflow["68"]["class_type"] = "EmptyLatentImage"

    # Inject your prompt into the correct CLIP Text Encode node ID (67)
    if "67" in prompt_workflow and "inputs" in prompt_workflow["67"]:
        prompt_workflow["67"]["inputs"]["text"] = user_prompt

    # Clear old remnants in both output and temp directories
    search_dirs = [os.path.join(COMFYUI_DIR, "output"), os.path.join(COMFYUI_DIR, "temp")]
    for d in search_dirs:
        if os.path.exists(d):
            for file in os.listdir(d):
                try:
                    os.remove(os.path.join(d, file))
                except Exception:
                    pass

    # Launch ComfyUI inside the GPU environment
    print("ZeroGPU allocated. Launching ComfyUI server instance...")
    comfy_process = subprocess.Popen(
        [sys.executable, os.path.join(COMFYUI_DIR, "main.py"), "--listen", "127.0.0.1", "--port", "8188", "--highvram"],
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True
    )

    # Poll local port until the server is alive
    server_ready = False
    for _ in range(25):
        time.sleep(1)
        try:
            with urllib.request.urlopen("http://127.0.0.1:8188/history", timeout=1) as r:
                if r.status == 200:
                    server_ready = True
                    break
        except Exception:
            continue

    if not server_ready:
        print("ComfyUI server failed to initialize within time constraints.")
        comfy_process.terminate()
        return None

    # Enqueue the workflow
    print("Server online. Enqueueing workflow...")
    p = {"prompt": prompt_workflow}
    data = json.dumps(p).encode('utf-8')
    req = urllib.request.Request("http://127.0.0.1:8188/prompt", data=data, headers={'Content-Type': 'application/json'})
    
    try:
        with urllib.request.urlopen(req) as response:
            res = json.loads(response.read().decode('utf-8'))
            print(f"Workflow running. Prompt ID: {res['prompt_id']}")
    except Exception as e:
        print(f"API execution dispatch failed: {e}")
        comfy_process.terminate()
        return None

    # Track output directory for the compiled image asset
    generated_image_path = None
    for _ in range(90):  
        time.sleep(1)
        for d in search_dirs:
            if os.path.exists(d):
                files = [os.path.join(d, f) for f in os.listdir(d) if os.path.isfile(os.path.join(d, f))]
                if files:
                    generated_image_path = max(files, key=os.path.getmtime)
                    break
        if generated_image_path:
            print(f"Asset generation complete: {generated_image_path}")
            break

    # Clean up the server process
    comfy_process.terminate()
    comfy_process.wait()
    
    return generated_image_path


# ==========================================
# 4. Gradio Web Interface Layout
# ==========================================
with gr.Blocks() as demo:
    gr.Markdown("# My Custom ComfyUI App")
    gr.Markdown("Enter a prompt below to run your custom ComfyUI workflow live via ZeroGPU allocations.")
    
    with gr.Row():
        with gr.Column():
            prompt_input = gr.Textbox(
                label="Prompt", 
                value="A 3D blocky rendering of a green creature in a tan robe and chest plate, dancing in a dedicated boombox setup. Bright colors, bold lines, blocky cel shading.",
                lines=5
            )
            submit_btn = gr.Button("Generate")
        with gr.Column():
            image_output = gr.Image(label="Result")
            
    submit_btn.click(
        fn=generate_image,
        inputs=prompt_input,
        outputs=image_output
    )

if __name__ == "__main__":
    demo.launch()