app.py

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

import random
import re
import numpy as np

import matplotlib.pyplot as plt
from matplotlib.backends.backend_pdf import PdfPages
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
    project_id: str
    chain_id: str
    session_id: 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, request: Request):
    # img_buffer = io.BytesIO()

    pattern = f"""```python([\s\S]*?)```"""
    codes = re.findall(pattern, code.code)

    if codes:
        print(codes)
        exec("\n".join(codes))
    else:
        print(code.code)
        exec(code.code)
    # img_buffer.seek(0)  # Reset buffer position

    file_path = f"graphs_{code.project_id}_{code.chain_id}_{code.session_id}.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(file_path) 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 {file_path}")

    
    # 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")
    return {
        "message": "Graph created succesfully!",
        "url": str(request.base_url) + f"chart/{code.project_id}/{code.chain_id}/{code.session_id}"
    }



@app.get("/chart/{project_id}/{chain_id}/{session_id}")
async def get_chart(project_id: str, chain_id: str, session_id: str):
    pdf_path = f"graphs_{project_id}_{chain_id}_{session_id}.pdf"
    return FileResponse(pdf_path, media_type="application/pdf")