import os
import subprocess
from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.responses import JSONResponse
from fastapi.staticfiles import StaticFiles
import uvicorn

UPLOAD_FOLDER = 'uploads_gradio'
OUTPUT_FOLDER = 'outputs_gradio'

# Ensure directories exist
os.makedirs(UPLOAD_FOLDER, exist_ok=True)
os.makedirs(OUTPUT_FOLDER, exist_ok=True)

def animate_image(file_path):
    """
    Process the uploaded image and generate an animated GIF.
    """
    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)

    # Validate extension
    allowed_extensions = ['.png', '.jpg', '.jpeg', '.bmp']
    if ext.lower() not in allowed_extensions:
        raise ValueError("Unsupported file type. Please upload png, jpg, jpeg, or bmp.")

    # Create output directory for this specific image
    char_anno_dir = os.path.join(OUTPUT_FOLDER, f"{base}_out")
    os.makedirs(char_anno_dir, exist_ok=True)

    # Run the image_to_animation script
    try:
        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 not os.path.exists(gif_path):
            raise FileNotFoundError("Failed to generate animation. Ensure the input image contains clear humanoid drawings.")
        return gif_path

    except subprocess.CalledProcessError as e:
        raise RuntimeError(f"Error during processing: {e}")
    except Exception as e:
        raise RuntimeError(f"Unexpected error: {e}")

app = FastAPI()

# Serve the outputs directory statically so that generated GIFs can be accessed via URL
app.mount("/outputs", StaticFiles(directory=OUTPUT_FOLDER), name="outputs")

@app.post("/animate")
async def animate_endpoint(file: UploadFile = File(...)):
    try:
        # Save uploaded file
        file_path = os.path.join(UPLOAD_FOLDER, file.filename)
        with open(file_path, "wb") as f:
            f.write(await file.read())

        gif_path = animate_image(file_path)
        # Construct a URL to the generated GIF
        relative_gif_dir = os.path.basename(os.path.dirname(gif_path))
        gif_filename = os.path.basename(gif_path)
        gif_url = f"/outputs/{relative_gif_dir}/{gif_filename}"

        return JSONResponse({"gif_url": gif_url})
    except ValueError as ve:
        raise HTTPException(status_code=400, detail=str(ve))
    except FileNotFoundError as fnfe:
        raise HTTPException(status_code=404, detail=str(fnfe))
    except RuntimeError as re:
        raise HTTPException(status_code=500, detail=str(re))
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Unexpected error: {e}")

if __name__ == "__main__":
    # Use the PORT environment variable provided by Hugging Face Spaces
    port = int(os.getenv("PORT", "7860"))
    uvicorn.run(app, host="0.0.0.0", port=port)