import gradio as gr
from fastapi import FastAPI, UploadFile, File
import uvicorn, uuid, os
from wan_runner import generate_video, generate_image

api = FastAPI()

@api.post("/generate_video")
async def api_generate_video(ref: UploadFile = File(...), first: UploadFile = File(...), last: UploadFile = File(...)):
    uid = uuid.uuid4().hex
    os.makedirs(uid, exist_ok=True)
    paths = [f"{uid}/{name}" for name in ["ref.png", "first.png", "last.png"]]
    for upload, path in zip([ref, first, last], paths):
        with open(path, "wb") as f:
            f.write(await upload.read())
    output = f"{uid}/output.mp4"
    generate_video(*paths, output)
    return {"video_path": output}

@api.post("/generate_image")
async def api_generate_image(ref: UploadFile = File(...), prompt: str = ""):
    uid = uuid.uuid4().hex
    os.makedirs(uid, exist_ok=True)
    ref_path = f"{uid}/ref.png"
    with open(ref_path, "wb") as f:
        f.write(await ref.read())
    output = f"{uid}/output.png"
    generate_image(ref_path, prompt, output)
    return {"image_path": output}

with gr.Blocks() as demo:
    with gr.Tab("動画生成"):
        ref_img = gr.Image(label="参照画像", type="pil")
        first_img = gr.Image(label="開始画像", type="pil")
        last_img = gr.Image(label="終了画像", type="pil")
        btn_video = gr.Button("動画を生成")
        output_video = gr.Video()

        def video_ui(ref, first, last):
            import tempfile
            with tempfile.TemporaryDirectory() as tmpdir:
                # Save uploaded images
                ref_path = f"{tmpdir}/ref.png"
                first_path = f"{tmpdir}/first.png"
                last_path = f"{tmpdir}/last.png"
                
                ref.save(ref_path)
                first.save(first_path)
                last.save(last_path)
                
                output = f"{uuid.uuid4().hex}.mp4"
                generate_video(ref_path, first_path, last_path, output)
                return output

        btn_video.click(video_ui, [ref_img, first_img, last_img], output_video)

    with gr.Tab("画像生成"):
        ref_img2 = gr.Image(label="参照画像", type="pil")
        prompt = gr.Textbox(label="画像プロンプト")
        btn_image = gr.Button("画像を生成")
        output_image = gr.Image()

        def image_ui(ref, prompt):
            import tempfile
            with tempfile.TemporaryDirectory() as tmpdir:
                # Save uploaded image
                ref_path = f"{tmpdir}/ref.png"
                ref.save(ref_path)
                
                output = f"{uuid.uuid4().hex}.png"
                generate_image(ref_path, prompt, output)
                return output

        btn_image.click(image_ui, [ref_img2, prompt], output_image)

app = gr.mount_gradio_app(api, demo, path="/")

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)