from typing import  Union

from fastapi import FastAPI, Request, status
from fastapi.responses import JSONResponse

from src.blackbox.blackbox_factory import BlackboxFactory
from fastapi.middleware.cors import CORSMiddleware
from injector import Injector
import contextvars
import asyncio

app = FastAPI()

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

injector = Injector()
blackbox_factory = injector.get(BlackboxFactory)

@app.post("/")
async def blackbox(blackbox_name: Union[str, None] = None, request: Request = None):
    if not blackbox_name:
        return await JSONResponse(content={"error": "blackbox_name is required"}, status_code=status.HTTP_400_BAD_REQUEST)
    try:
        box = blackbox_factory.get_blackbox(blackbox_name)
    except ValueError:
        return await JSONResponse(content={"error": "value error"}, status_code=status.HTTP_400_BAD_REQUEST)
    return await box.fast_api_handler(request)

user_id_var = contextvars.ContextVar('user_id')

async def perform_action():
    user_id1 = user_id_var.get()
    await asyncio.sleep(5)
    user_id = user_id_var.get()
    # Perform action based on user ID
    print(f"{user_id1}:{user_id}")

@app.post("/llm")
async def llm(request: Request = None):
    
    user_id_var.set(request.headers["context-id"])

    print(request.headers["context-id"])
    # await perform_action()