app.py

from fastapi import FastAPI, Request, Response
from fastapi.responses import HTMLResponse, RedirectResponse
from fastapi.templating import Jinja2Templates
from google.oauth2 import id_token
from google.auth.transport import requests as google_requests
from google_auth_oauthlib.flow import Flow
from chainlit.utils import mount_chainlit
import secrets
import json
import base64
from constants import (
    OAUTH_GOOGLE_CLIENT_ID,
    OAUTH_GOOGLE_CLIENT_SECRET,
    CHAINLIT_URL,
)
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
import os


GOOGLE_CLIENT_ID = OAUTH_GOOGLE_CLIENT_ID
GOOGLE_CLIENT_SECRET = OAUTH_GOOGLE_CLIENT_SECRET
GOOGLE_REDIRECT_URI = f"{CHAINLIT_URL}/auth/oauth/google/callback"

app = FastAPI()
# app.mount("/public", StaticFiles(directory="public"), name="public")
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Update with appropriate origins
    allow_methods=["*"],
    allow_headers=["*"],  # or specify the headers you want to allow
    expose_headers=["X-User-Info"],  # Expose the custom header
)

templates = Jinja2Templates(directory="templates")
session_store = {}
CHAINLIT_PATH = "/chainlit_tutor"

USER_ROLES = {
    "tgardos@bu.edu": ["instructor", "bu"],
    "xthomas@bu.edu": ["instructor", "bu"],
    "faridkar@bu.edu": ["instructor", "bu"],
    "xavierohan1@gmail.com": ["guest"],
    # Add more users and roles as needed
}

# Create a Google OAuth flow
flow = Flow.from_client_config(
    {
        "web": {
            "client_id": GOOGLE_CLIENT_ID,
            "client_secret": GOOGLE_CLIENT_SECRET,
            "auth_uri": "https://accounts.google.com/o/oauth2/auth",
            "token_uri": "https://oauth2.googleapis.com/token",
            "redirect_uris": [GOOGLE_REDIRECT_URI],
            "scopes": [
                "openid",
                # "https://www.googleapis.com/auth/userinfo.email",
                # "https://www.googleapis.com/auth/userinfo.profile",
            ],
        }
    },
    scopes=[
        "openid",
        "https://www.googleapis.com/auth/userinfo.email",
        "https://www.googleapis.com/auth/userinfo.profile",
    ],
    redirect_uri=GOOGLE_REDIRECT_URI,
)


def get_user_role(username: str):
    return USER_ROLES.get(username, ["student"])  # Default to "student" role


def get_user_info_from_cookie(request: Request):
    user_info_encoded = request.cookies.get("X-User-Info")
    if user_info_encoded:
        try:
            user_info_json = base64.b64decode(user_info_encoded).decode()
            return json.loads(user_info_json)
        except Exception as e:
            print(f"Error decoding user info: {e}")
            return None
    return None


def get_user_info(request: Request):
    session_token = request.cookies.get("session_token")
    if session_token and session_token in session_store:
        return session_store[session_token]
    return None


@app.get("/", response_class=HTMLResponse)
async def login_page(request: Request):
    user_info = get_user_info_from_cookie(request)
    if user_info and user_info.get("google_signed_in"):
        return RedirectResponse("/post-signin")
    return templates.TemplateResponse("login.html", {"request": request})


@app.get("/login/guest")
@app.post("/login/guest")
async def login_guest():
    username = "guest"
    session_token = secrets.token_hex(16)
    unique_session_id = secrets.token_hex(8)
    username = f"{username}_{unique_session_id}"
    session_store[session_token] = {
        "email": username,
        "name": "Guest",
        "profile_image": "",
        "google_signed_in": False,  # Ensure guest users do not have this flag
    }
    user_info_json = json.dumps(session_store[session_token])
    user_info_encoded = base64.b64encode(user_info_json.encode()).decode()

    # Set cookies
    response = RedirectResponse(url="/post-signin", status_code=303)
    response.set_cookie(key="session_token", value=session_token)
    response.set_cookie(key="X-User-Info", value=user_info_encoded, httponly=True)
    return response


@app.get("/login/google")
async def login_google(request: Request):
    # Clear any existing session cookies to avoid conflicts with guest sessions
    response = RedirectResponse(url="/post-signin")
    response.delete_cookie(key="session_token")
    response.delete_cookie(key="X-User-Info")

    user_info = get_user_info_from_cookie(request)
    print(f"User info: {user_info}")
    # Check if user is already signed in using Google
    if user_info and user_info.get("google_signed_in"):
        return RedirectResponse("/post-signin")
    else:
        authorization_url, _ = flow.authorization_url(prompt="consent")
        return RedirectResponse(authorization_url, headers=response.headers)


@app.get("/auth/oauth/google/callback")
async def auth_google(request: Request):
    try:
        flow.fetch_token(code=request.query_params.get("code"))
        credentials = flow.credentials
        user_info = id_token.verify_oauth2_token(
            credentials.id_token, google_requests.Request(), GOOGLE_CLIENT_ID
        )

        email = user_info["email"]
        name = user_info.get("name", "")
        profile_image = user_info.get("picture", "")

        session_token = secrets.token_hex(16)
        session_store[session_token] = {
            "email": email,
            "name": name,
            "profile_image": profile_image,
            "google_signed_in": True,  # Set this flag to True for Google-signed users
        }

        user_info_json = json.dumps(session_store[session_token])
        user_info_encoded = base64.b64encode(user_info_json.encode()).decode()

        # Set cookies
        response = RedirectResponse(url="/post-signin", status_code=303)
        response.set_cookie(key="session_token", value=session_token)
        response.set_cookie(key="X-User-Info", value=user_info_encoded, httponly=True)
        return response
    except Exception as e:
        print(f"Error during Google OAuth callback: {e}")
        return RedirectResponse(url="/", status_code=302)


@app.get("/post-signin", response_class=HTMLResponse)
async def post_signin(request: Request):
    user_info = get_user_info_from_cookie(request)
    if not user_info:
        user_info = get_user_info(request)
    # if user_info and user_info.get("google_signed_in"):
    if user_info:
        username = user_info["email"]
        role = get_user_role(username)
        jwt_token = request.cookies.get("X-User-Info")
        return templates.TemplateResponse(
            "dashboard.html",
            {
                "request": request,
                "username": username,
                "role": role,
                "jwt_token": jwt_token,
            },
        )
    return RedirectResponse("/")


@app.post("/start-tutor")
async def start_tutor(request: Request):
    user_info = get_user_info_from_cookie(request)
    if user_info:
        user_info_json = json.dumps(user_info)
        user_info_encoded = base64.b64encode(user_info_json.encode()).decode()

        response = RedirectResponse(CHAINLIT_PATH, status_code=303)
        response.set_cookie(key="X-User-Info", value=user_info_encoded, httponly=True)
        return response

    return RedirectResponse(url="/")


@app.exception_handler(Exception)
async def exception_handler(request: Request, exc: Exception):
    return templates.TemplateResponse(
        "error.html", {"request": request, "error": str(exc)}, status_code=500
    )


@app.get("/chainlit_tutor/logout", response_class=HTMLResponse)
@app.post("/chainlit_tutor/logout", response_class=HTMLResponse)
async def app_logout(request: Request, response: Response):
    # Clear session cookies
    response.delete_cookie("session_token")
    response.delete_cookie("X-User-Info")

    print("logout_page called")

    # Redirect to the logout page with embedded JavaScript
    return RedirectResponse(url="/", status_code=302)


mount_chainlit(app=app, target="main.py", path=CHAINLIT_PATH)

if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="127.0.0.1", port=7860)