import os
import subprocess
from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.responses import FileResponse
from fastapi.middleware.cors import CORSMiddleware
import uvicorn

# Define directories for uploads and outputs
UPLOAD_FOLDER = 'uploads_gradio'
OUTPUT_FOLDER = 'outputs_gradio'

# Create directories if they don't exist
os.makedirs(UPLOAD_FOLDER, exist_ok=True)
os.makedirs(OUTPUT_FOLDER, exist_ok=True)

app = FastAPI(
    title="Animated Drawings API",
    description="Upload your drawing and receive an animated GIF.",
    version="1.0.0"
)

# Enable CORS (optional, adjust origins as needed)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Update with specific origins in production
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

def animate_image(file_path: str) -> str:
    """
    Process the uploaded image and generate an animated GIF.

    Args:
        file_path (str): Path to the uploaded file.

    Returns:
        str: Path to the generated 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)

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

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

        # Run the image_to_animation.py script with required arguments
        subprocess.run([
            'python', 'examples/image_to_animation.py',
            input_path, char_anno_dir
        ], check=True)

        # Path to the generated GIF
        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. Please ensure the input image contains clear humanoid drawings.")

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

@app.post("/animate", summary="Generate Animated GIF from Image")
async def generate_gif(image: UploadFile = File(...)):
    """
    Endpoint to upload an image and receive an animated GIF.

    Args:
        image (UploadFile): The image file to be animated.

    Returns:
        FileResponse: The generated animated GIF.
    """
    # Validate the uploaded file
    if not image:
        raise HTTPException(status_code=400, detail="No file uploaded.")

    filename = image.filename
    base, ext = os.path.splitext(filename)
    allowed_extensions = ['.png', '.jpg', '.jpeg', '.bmp']
    if ext.lower() not in allowed_extensions:
        raise HTTPException(status_code=400, detail="Unsupported file type. Please upload an image file (png, jpg, jpeg, bmp).")

    # Save the uploaded file to the upload directory
    upload_path = os.path.join(UPLOAD_FOLDER, filename)
    try:
        with open(upload_path, "wb") as f:
            f.write(await image.read())
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Failed to save uploaded file: {e}")

    # Process the image to generate GIF
    try:
        gif_path = animate_image(upload_path)
    except ValueError as ve:
        raise HTTPException(status_code=400, detail=str(ve))
    except FileNotFoundError as fnfe:
        raise HTTPException(status_code=500, 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"An unexpected error occurred: {e}")

    # Return the generated GIF as a response
    if os.path.exists(gif_path):
        return FileResponse(path=gif_path, media_type="image/gif", filename="animated.gif")
    else:
        raise HTTPException(status_code=500, detail="Failed to generate GIF.")

@app.get("/", summary="Root Endpoint")
def read_root():
    return {"message": "Welcome to the Animated Drawings API. Use the /animate endpoint to upload images and receive animated GIFs."}

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