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

import random
import re

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")