import gradio as gr
import os
import subprocess
import requests
import time

UPLOAD_FOLDER = 'uploads_gradio'
OUTPUT_FOLDER = 'outputs_gradio'

os.makedirs(UPLOAD_FOLDER, exist_ok=True)
os.makedirs(OUTPUT_FOLDER, exist_ok=True)

# Start TorchServe in the background
# Adjust arguments to match your configuration
ts_process = subprocess.Popen([
    "torchserve", "--start",
    "--model-store", "/home/torchserve/model-store",
    "--ts-config", "/home/torchserve/config.properties",
    "--no-config-snapshots"
])

# Wait for TorchServe to become healthy
while True:
    try:
        r = requests.get("http://127.0.0.1:8080/ping")
        if r.status_code == 200:
            break
    except:
        pass
    time.sleep(1)

def animate_image(file_path):
    # Here you can call TorchServe endpoints if needed, for example:
    # response = requests.post("http://127.0.0.1:8080/predictions/drawn_humanoid_detector", files={"data": open(file_path, "rb")})
    # Process the response and then run your animation logic.
    # ...
    # Or if you run your original script, just as before:
    
    if not file_path:
        raise ValueError("No file uploaded.")
    input_path = file_path
    filename = os.path.basename(input_path)
    base, ext = os.path.splitext(filename)
    allowed_extensions = ['.png', '.jpg', '.jpeg', '.bmp']
    if ext.lower() not in allowed_extensions:
        raise ValueError("Unsupported file type.")

    char_anno_dir = os.path.join(OUTPUT_FOLDER, f"{base}_out")
    os.makedirs(char_anno_dir, exist_ok=True)

    subprocess.run(['python', 'examples/image_to_animation.py', input_path, char_anno_dir], check=True)

    gif_path = os.path.join(char_anno_dir, "video.gif")
    if os.path.exists(gif_path):
        return gif_path
    else:
        raise FileNotFoundError("Animation failed to generate.")

iface = gr.Interface(
    fn=animate_image,
    inputs=gr.Image(label="Upload Drawing", type="filepath", sources=["upload", "webcam"]),
    outputs=gr.Image(label="Animated GIF"),
    title="Animated Drawings",
    description="Upload or take a photo of a drawing, get an animated GIF."
)

if __name__ == "__main__":
    # Gradio on port 7860
    iface.launch(server_name="0.0.0.0", server_port=7860)