app/main.py

import os
import chromadb

from fastapi import FastAPI, Request, Form, File, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from fastapi.responses import HTMLResponse, RedirectResponse
from fastapi.templating import Jinja2Templates

from admin import admin_functions as admin
from utils.db import UserFaceEmbeddingFunction,ChromaDBFaceHelper
from api import userlogin, userlogout, userchat, userupload
from utils.db import ChromaDBFaceHelper
from utils.chat_rag import LlamaModelSingleton

app = FastAPI()

CHROMADB_LOC = os.getenv('CHROMADB_LOC')

# Add middleware
# Set all origins to wildcard for simplicity, but we should limit this in production
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Persitent storage for chromadb setup in /data volume
ec_client = chromadb.PersistentClient(CHROMADB_LOC)

# The following collection reference is needed for admin function to register face
user_faces_db = ec_client.get_or_create_collection(name="user_faces_db", embedding_function=UserFaceEmbeddingFunction())


@app.on_event("startup")
async def startup_event():
    global chromadb_face_helper
    # Assuming chromadb persistent store client for APIs is in helper
    db_path = CHROMADB_LOC
    chromadb_face_helper = ChromaDBFaceHelper(db_path) # Used by APIs
    
    # Perform any other startup tasks here
    # Preload the LLM model
    await LlamaModelSingleton.get_instance()
    print("LLM model loaded and ready.")

    print(f"MODEL_PATH in main.py = {os.getenv('MODEL_PATH')} ")

# Mount static files
app.mount("/static", StaticFiles(directory="static"), name="static")

# Setup Jinja2Templates to point to the templates directory
templates = Jinja2Templates(directory="admin/templates")

@app.get("/")
async def get_admin_login(request: Request):
    return templates.TemplateResponse("admin_login.html", {"request": request})

# Admin Login Handler
@app.post("/admin/login", response_class=HTMLResponse)
async def handle_admin_login(request: Request, username: str = Form(...), password: str = Form(...)):

    if admin.verify_admin_password(username, password):
        # Redirect to user registration page upon successful login
        return RedirectResponse(url="/admin/register_user", status_code=303)
    else:
        # Reload login page with error message
        return templates.TemplateResponse("admin_login.html", {"request": request, "error": "Invalid password"})

# To display the register user page
@app.get("/admin/register_user", response_class=HTMLResponse)
async def get_user_registration(request: Request):
    # Render the registration form
    return templates.TemplateResponse("user_registration.html", {"request": request})

# User Registration Handler
@app.post("/admin/register_user", response_class=HTMLResponse)
async def handle_user_registration(request: Request, email: str = Form(...), name: str = Form(...), role: str = Form(...), file: UploadFile = File(...)):
    user_id = await admin.register_user(user_faces_db, email, name, role, file)
    if user_id:
        # Calculate disk usage
        disk_usage = admin.get_disk_usage("/home/user/data")

        # Redirect or display a success message
        return templates.TemplateResponse("registration_success.html", {
            "request": request,
            "disk_usage": disk_usage 
        })
    else:
        # Reload registration page with error message
        return templates.TemplateResponse("user_registration.html", {"request": request, "error": "Registration failed"})

# To display admin utilities
@app.get("/admin/data_management", response_class=HTMLResponse)
async def get_db_details(request: Request):
    # Render the Chroma DB details
    faces = admin.faces_count(ec_client, user_faces_db)
    return templates.TemplateResponse("data_management.html", {
        "request": request,
        "faces" : faces
    })

@app.post("/admin/delete_faces")
async def delete_faces(request: Request):
    try:
        # Call your function to remove all faces
        admin.remove_all_faces(ec_client)
        request.session['flash'] = "All user data successfully deleted."
    except Exception as e:
        request.session['flash'] = f"Failed to delete user data: {str(e)}"
    
    return RedirectResponse(url="/admin/data_management", status_code=303)

app.include_router(userlogin.router)
app.include_router(userlogout.router)
app.include_router(userchat.router)
app.include_router(userupload.router)