from fastapi import FastAPI, HTTPException
from fastapi.responses import JSONResponse
from argopy import DataFetcher
import matplotlib.pyplot as plt
import io
import base64

app = FastAPI()

@app.get("/plot_trajectory")
def plot_trajectory(lon_min: float, lon_max: float, lat_min: float, lat_max: float, 
                    depth_min: float, depth_max: float, 
                    date_start: str, date_end: str):
    try:
        # Create data fetcher with user inputs
        f = DataFetcher().region([lon_min, lon_max, lat_min, lat_max, depth_min, depth_max, date_start, date_end])
        
        # Plotting
        plt.figure()
        f.plot('trajectory', add_legend=False)

        # Save the plot to a BytesIO buffer
        buf = io.BytesIO()
        plt.savefig(buf, format='png')
        buf.seek(0)
        plt.close()

        # Encode the image in base64 and prepare it for JSON response
        encoded_img = base64.b64encode(buf.read()).decode('utf-8')
        response = {"image_base64": encoded_img}

        return JSONResponse(content=response)

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


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