app/main.py

from fastapi import FastAPI, Request
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse, RedirectResponse
from contextlib import asynccontextmanager
from app.routers import shield, auth
from app.models.deberta_model import classifier
from app.database.db import init_db
import os

@asynccontextmanager
async def lifespan(app: FastAPI):
    # Startup: Initialize Database & Load classifier model
    print("Initializing Database...")
    try:
        init_db()
        print("Database initialized successfully!")
    except Exception as e:
        print(f"Error initializing DB: {e}")

    print("Loading DeBERTa-v3 model. This might take a few moments on the first run...")
    try:
        classifier.initialize()
        print("DeBERTa-v3 model loaded successfully!")
    except Exception as e:
        print(f"Error pre-loading model: {e}")

    print("Initializing SemanticGuard embedding model...")
    try:
        from app.services.semantic_guard import semantic_guard
        semantic_guard.initialize()
        print("SemanticGuard embedding model loaded successfully!")
    except Exception as e:
        print(f"Error pre-loading SemanticGuard: {e}")

    yield
    # Shutdown: Clean up if needed
    pass

app = FastAPI(
    title="GuardrailShield AI",
    description="Defensive Guardrail System sitting in front of LLMs",
    version="1.0.0",
    lifespan=lifespan
)

# Include routers
app.include_router(shield.router)
app.include_router(auth.router)

# Mount Static Files
static_dir = os.path.join(os.path.dirname(__file__), "static")
if os.path.exists(static_dir):
    app.mount("/static", StaticFiles(directory=static_dir), name="static")

@app.get("/")
def get_chat_client():
    """Serves the public user chat interface."""
    return FileResponse(os.path.join(static_dir, "chat.html"))

@app.get("/login")
def get_login_page():
    """Serves the admin login page."""
    return FileResponse(os.path.join(static_dir, "login.html"))

@app.get("/dashboard")
def get_dashboard_ui(request: Request):
    """Serves the security admin dashboard."""
    return FileResponse(os.path.join(static_dir, "index.html"))

@app.get("/health", tags=["Health"])
def health_check():
    return {"status": "healthy"}