import gradio as gr
import os
import uuid
import subprocess
import time
from threading import Thread
from queue import Queue
from prompt_builder import enhance_prompt

job_queue = Queue()
results = {}

OUTPUT_DIR = "/workspace/output"
os.makedirs(OUTPUT_DIR, exist_ok=True)


# Worker
def worker():
    while True:
        job_id, prompt, mode = job_queue.get()
        try:
            enhanced = enhance_prompt(prompt)

            prompt_path = f"/tmp/{job_id}.txt"
            with open(prompt_path, "w") as f:
                f.write(enhanced)

            # Select pipeline
            script = {
                "base": "base/run.sh",
                "distill": "distill/run.sh",
                "sr540": "sr_540p/run.sh",
                "sr1080": "sr_1080p/run.sh"
            }[mode]

            cmd = f"bash MagiCompiler/example/{script}"
            process = subprocess.Popen(cmd, shell=True)

            # STREAM frames
            for i in range(30):
                fake_frame = f"/tmp/frame_{i}.png"
                with open(fake_frame, "wb") as f:
                    f.write(b"")

                results[job_id] = {
                    "frame": fake_frame,
                    "status": "running"
                }
                time.sleep(0.5)

            process.wait()

            # Final outputs (replace with real paths)
            video_path = f"{OUTPUT_DIR}/{job_id}.mp4"
            audio_path = f"{OUTPUT_DIR}/{job_id}.wav"

            # fake output
            open(video_path, "wb").close()
            open(audio_path, "wb").close()

            results[job_id] = {
                "video": video_path,
                "audio": audio_path,
                "status": "done"
            }

        except Exception as e:
            results[job_id] = {"error": str(e)}

        job_queue.task_done()


Thread(target=worker, daemon=True).start()


def submit(prompt, mode):
    job_id = str(uuid.uuid4())
    job_queue.put((job_id, prompt, mode))
    return job_id


def generate(prompt, mode):
    job_id = submit(prompt, mode)

    while True:
        if job_id in results:
            res = results[job_id]

            if res.get("status") == "running":
                yield res["frame"], None, None

            elif res.get("status") == "done":
                yield None, res["video"], res["audio"]
                break

            elif "error" in res:
                raise RuntimeError(res["error"])

        time.sleep(1)


# UI
with gr.Blocks() as demo:
    gr.Markdown("# 🎬 MagiHuman AI Generator")

    prompt = gr.Textbox(label="Enter your prompt")
    mode = gr.Dropdown(["base", "distill", "sr540", "sr1080"], value="base")

    btn = gr.Button("Generate")

    frame = gr.Image(label="Live Preview")
    video = gr.Video(label="Final Video")
    audio = gr.Audio(label="Generated Audio")

    btn.click(generate, inputs=[prompt, mode], outputs=[frame, video, audio])

demo.queue().launch(server_name="0.0.0.0", server_port=7860)