app.py

from fastapi import FastAPI
from fastapi.responses import Response
from fastapi.responses import FileResponse
from pydantic import BaseModel

import random

import matplotlib.pyplot as plt
import pandas as pd

import io

app = FastAPI()

# Define the ticket schema using Pydantic
class Ticket(BaseModel):
    name: str
    department: str
    category: str
    description: str
    service_category: str
    difficulty: int  # Adjust type as needed (e.g., int or str)


class Code(BaseModel):
    code: str

@app.get("/")
def greet_json():
    return {"Hello": "World!"}

@app.post("/ticket")
async def create_ticket(ticket: Ticket):
    # Here you can process the ticket, e.g., save it to a database.
    # For now, we simply return the received ticket data.
    tick = ticket.dict()
    tick["number"] = random.randint(1000, 9999)
    return {
        "message": "Ticket created successfully",
        "ticket": tick
    }


@app.post("/run_code")
async def run_code(code: Code):
    # img_buffer = io.BytesIO()
    exec(code.code)
    # img_buffer.seek(0)  # Reset buffer position

    file_path = "graph.pdf"

    # if "plt.subplots(" in code.code:
    #     print("SUBPLOTS DETECTED")
    #     exec(code.code + "\nfig.savefig(file_path)\nplt.close()")
    # else:
    #     print("NO SUBPLOTS")
    #     exec(code.code + "\nplt.savefig(file_path)\nplt.close()")

    # plt.close()
    # plt.savefig(file_path)
    # plt.close()

    # Get all open figures
    figures = [plt.figure(i) for i in plt.get_fignums()]
    
    # Save all figures in a single PDF using PdfPages
    pdf_filename = "all_graphs.pdf"
    with PdfPages(pdf_filename) as pdf:
        for fig in figures:
            pdf.savefig(fig)  # Save each figure as a page in the PDF
            plt.close(fig)  # Close the figure to free memory
    
    print(f"Saved all figures to {pdf_filename}")

    
    # return Response(content=img_buffer.getvalue(), media_type="image/png")
    # return FileResponse(file_path, media_type="image/png")
    return FileResponse(file_path, media_type="application/pdf", filename="graph.pdf")