main.py

from collections import defaultdict
from fastapi import (
    APIRouter,
    FastAPI,
    status,
    Depends,
    HTTPException,
    BackgroundTasks,
    Response,
)
from sqlalchemy.orm import Session
from fastapi.middleware.cors import CORSMiddleware
from fastapi.websockets import WebSocket
from fastapi.responses import StreamingResponse
from contextlib import asynccontextmanager
import json
from jose import jwt
from jose import ExpiredSignatureError
from io import StringIO
import csv

import schemas
import crud
import models
from database import Base, engine, get_db, SessionLocal
from utils import check_user_level
import config
from security import jwt_cookie, get_jwt_user
from routes.tests import testRouter
from routes.studies import studiesRouter
from routes.chat import chatRouter

websocket_users = defaultdict(lambda: defaultdict(set))
websocket_users_global = defaultdict(set)


@asynccontextmanager
async def lifespan(app: FastAPI):
    db = SessionLocal()
    try:
        if crud.get_user_by_email(db, config.ADMIN_EMAIL) is None:
            crud.create_user(
                db,
                schemas.UserCreate(
                    email=config.ADMIN_EMAIL,
                    nickname=config.ADMIN_NICKNAME,
                    password=config.ADMIN_PASSWORD,
                    type=models.UserType.ADMIN.value,
                ),
            )
    finally:
        db.close()
    yield


Base.metadata.create_all(bind=engine)

app = FastAPI(lifespan=lifespan)

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

apiRouter = APIRouter(prefix="/tmp-api")
v1Router = APIRouter(prefix="/v1")
authRouter = APIRouter(prefix="/auth", tags=["auth"])
usersRouter = APIRouter(prefix="/users", tags=["users"])
sessionsRouter = APIRouter(prefix="/sessions", tags=["sessions"])
studyRouter = APIRouter(prefix="/studies", tags=["studies"])
websocketRouter = APIRouter(prefix="/ws", tags=["websocket"])


@v1Router.get("/health", status_code=status.HTTP_204_NO_CONTENT)
def health():
    return


@authRouter.post("/login", status_code=200)
def login(
    login: schemas.LoginData,
    response: Response,
    db: Session = Depends(get_db),
):
    db_user = crud.get_user_by_email_and_password(db, login.email, login.password)
    if db_user is None:
        raise HTTPException(status_code=401, detail="Incorrect email or password")

    subject = {"uid": db_user.id, "email": db_user.email, "nickname": db_user.nickname}
    access_token = jwt_cookie.create_access_token(subject)

    jwt_cookie.set_access_cookie(
        response=response,
        access_token=access_token,
        expires_delta=jwt_cookie.access_expires_delta,
    )

    refresh_token = jwt_cookie.create_refresh_token(subject)

    jwt_cookie.set_refresh_cookie(
        response=response,
        refresh_token=refresh_token,
        expires_delta=jwt_cookie.refresh_expires_delta,
    )


@authRouter.post("/register", status_code=status.HTTP_201_CREATED)
def register(
    register: schemas.RegisterData,
    db: Session = Depends(get_db),
):
    db_user = crud.get_user_by_email(db, email=register.email)
    if db_user:
        raise HTTPException(status_code=400, detail="User already registered")

    user_data = schemas.UserCreate(
        email=register.email,
        password=register.password,
        nickname=register.nickname,
        type=(
            models.UserType.TUTOR.value
            if register.is_tutor
            else models.UserType.STUDENT.value
        ),
    )

    user = crud.create_user(db=db, user=user_data)

    return user.id


@usersRouter.post("", status_code=status.HTTP_201_CREATED)
def create_user(
    user: schemas.UserCreate,
    db: Session = Depends(get_db),
    current_user: schemas.User = Depends(get_jwt_user),
):
    if not check_user_level(current_user, models.UserType.ADMIN):
        raise HTTPException(
            status_code=401, detail="You do not have permission to create a user"
        )

    db_user = crud.get_user_by_email(db, email=user.email)
    if db_user:
        raise HTTPException(status_code=400, detail="User already registered")

    user = crud.create_user(db=db, user=user)

    return user.id


@usersRouter.get("/{user_id}", response_model=schemas.User)
def read_user(
    user_id: int,
    db: Session = Depends(get_db),
    current_user: schemas.User = Depends(get_jwt_user),
):
    if (
        not check_user_level(current_user, models.UserType.ADMIN)
        and current_user.id != user_id
    ):
        raise HTTPException(
            status_code=401, detail="You do not have permission to view this user"
        )

    db_user = crud.get_user(db, user_id=user_id)
    if db_user is None:
        raise HTTPException(status_code=404, detail="User not found")

    return db_user


@usersRouter.get("", response_model=list[schemas.User])
def read_users(
    skip: int = 0,
    db: Session = Depends(get_db),
    current_user: schemas.User = Depends(get_jwt_user),
):
    return crud.get_users(db, skip=skip)


@usersRouter.delete("/{user_id}", status_code=status.HTTP_204_NO_CONTENT)
def delete_user(
    user_id: int,
    db: Session = Depends(get_db),
    current_user: schemas.User = Depends(get_jwt_user),
):
    if not check_user_level(current_user, models.UserType.ADMIN):
        raise HTTPException(
            status_code=401, detail="You do not have permission to delete a user"
        )

    if not crud.delete_user(db, user_id):
        raise HTTPException(status_code=404, detail="User not found")


@usersRouter.patch("/{user_id}", status_code=status.HTTP_204_NO_CONTENT)