from fastapi import FastAPI, HTTPException, Request, Query
from fastapi.middleware.cors import CORSMiddleware
from fastapi_cache import FastAPICache
from fastapi_cache.backends.inmemory import InMemoryBackend
from fastapi_cache.decorator import cache

app = FastAPI()

@app.on_event("startup")
async def startup():
    FastAPICache.init(InMemoryBackend(), prefix="fastapi-cache")

from pydantic import BaseModel
from typing import List, Dict, Any

TOGETHER_API_KEY = os.getenv('TOGETHER_API_KEY')
BRAVE_API_KEY = os.getenv('BRAVE_API_KEY')
GROQ_API_KEY = os.getenv("GROQ_API_KEY")
HELICON_API_KEY = os.getenv("HELICON_API_KEY")
SUPABASE_USER = os.environ['SUPABASE_USER']
SUPABASE_PASSWORD = os.environ['SUPABASE_PASSWORD']

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],)

import asyncio
import uvicorn
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from jupyter_client import KernelManager
from typing import Dict
from datetime import datetime, timedelta
import psutil

app = FastAPI()

# Model for input data
class CodeExecutionRequest(BaseModel):
    session_token: str
    code: str

# Store kernel managers and last access times
kernel_managers: Dict[str, KernelManager] = {}
last_access_times: Dict[str, datetime] = {}

# Timeout duration in seconds
TIMEOUT_DURATION = 600  # 10 minutes

# Function to create a new kernel
async def create_kernel(session_token: str):
    km = KernelManager()
    km.start_kernel()
    kernel_managers[session_token] = km
    last_access_times[session_token] = datetime.now()

# Function to kill a kernel
async def kill_kernel(session_token: str):
    km = kernel_managers.pop(session_token, None)
    if km:
        km.shutdown_kernel(now=True)
    last_access_times.pop(session_token, None)

# Function to execute code in a kernel
async def execute_code(session_token: str, code: str):
    if session_token not in kernel_managers:
        await create_kernel(session_token)

    km = kernel_managers[session_token]
    kc = km.client()
    kc.start_channels()

    try:
        kc.execute(code)
        reply = kc.get_shell_msg(timeout=5)
        result = reply['content']
    except Exception as e:
        result = {'status': 'error', 'ename': type(e).__name__, 'evalue': str(e)}

    kc.stop_channels()
    last_access_times[session_token] = datetime.now()
    return result

# Background task to check for idle kernels
async def check_idle_kernels():
    while True:
        now = datetime.now()
        for session_token, last_access in list(last_access_times.items()):
            if now - last_access > timedelta(seconds=TIMEOUT_DURATION):
                await kill_kernel(session_token)
        await asyncio.sleep(60)  # Check every minute

@app.on_event("startup")
async def startup_event():
    asyncio.create_task(check_idle_kernels())

@app.post("/execute")
async def execute(request: CodeExecutionRequest):
    result = await execute_code(request.session_token, request.code)
    return result

@app.get("/info")
async def get_info():
    # Get the number of active kernels
    active_kernels = len(kernel_managers)

    # Get system resource usage
    cpu_usage = psutil.cpu_percent(interval=1)
    ram_usage = psutil.virtual_memory().percent

    # Return the information
    return {
        "active_kernels": active_kernels,
        "cpu_usage_percent": cpu_usage,
        "ram_usage_percent": ram_usage
    }

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