from fastapi import FastAPI, Response
import papermill as pm
import io

app = FastAPI()

@app.get("/")
def run_notebook(fruits: str):
    fruit_list = fruits.split(",")

    # Execute the notebook with Papermill
    output_path = "output.ipynb"
    pm.execute_notebook(
        "fruitchecker.ipynb",
        output_path,
        parameters={"fruits": fruit_list}
    )

    # Extract the CSV result
    with open(output_path, "r") as f:
        notebook_content = f.read()
    
    # Find the CSV output in the notebook
    csv_output = notebook_content.split("csv_output = ")[-1].strip().strip("\"'")

    # Return as downloadable CSV
    return Response(content=csv_output, media_type="text/csv",
                    headers={"Content-Disposition": "attachment; filename=fruits.csv"})