from fastapi import FastAPI, Request, UploadFile, File
from fastapi.responses import StreamingResponse
from fastapi.responses import JSONResponse
from rembg import remove
from PIL import Image
import datetime
import io

app = FastAPI()
app.router.redirect_slashes = False

@app.api_route("/", methods=["GET", "HEAD"])
async def index(request: Request):
    return JSONResponse({
        "status": "ok",
        "message": "Server is alive",
        "timestamp": datetime.datetime.utcnow().isoformat() + "Z"
    })


@app.post('/remove-bg')
async def remove_bg(image: UploadFile = File(...)):
    """Accepts an uploaded image file, removes its background and returns a PNG.

    Input: multipart form file field named `image`.
    Output: PNG image stream with Content-Disposition attachment `no-bg.png`.
    """
    contents = await image.read()
    input_image = Image.open(io.BytesIO(contents))

    # perform background removal (rembg accepts a PIL Image)
    output = remove(input_image)

    buffer = io.BytesIO()
    output.save(buffer, format='PNG')
    buffer.seek(0)

    headers = {"Content-Disposition": "attachment; filename=\"no-bg.png\""}
    return StreamingResponse(buffer, media_type='image/png', headers=headers)

# If you want to run locally with `python api.py` for quick testing
if __name__ == '__main__':
    import uvicorn
    uvicorn.run(app, host='0.0.0.0', port=7860)