main.py

from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Depends, HTTPException, status, Body, UploadFile, File, Query
from fastapi.responses import HTMLResponse, StreamingResponse, JSONResponse
from fastapi.security import OAuth2PasswordRequestForm
from sqlalchemy.orm import Session
from typing import List, Dict, Any, Optional
import json
import asyncio
from datetime import datetime, timedelta
import uuid
from io import BytesIO
import time
import os
import base64
import shutil
import tempfile
import zipfile

from server import models, schemas, auth
from server import database
from server.database import get_db

# NEW: Import services and routers
from server.services.cleanup_service import perform_startup_cleanup
from server.websocket_manager import websocket_endpoint # Import the websocket_endpoint
from server.routers import auth_router, agent_router, alert_router, server_router

app = FastAPI()

@app.on_event("startup")
async def on_startup():
    database.Base.metadata.create_all(bind=database.engine)
    await perform_startup_cleanup() # Run cleanup on startup

@app.get("/")
async def get():
    return HTMLResponse("<h1>Remote Admin & Monitoring Server</h1>")

# --- Network Test Endpoints (only ping remains) ---
@app.get("/network-test/ping")
async def network_test_ping():
    """A simple endpoint to measure latency."""
    return {"message": "pong", "timestamp": datetime.now().isoformat()}
# --- End Network Test Endpoints ---

# NEW: Include routers
app.include_router(auth_router.router, tags=["Authentication"])
app.include_router(agent_router.router, tags=["Agents"])
app.include_router(alert_router.router, tags=["Alerts"])
app.include_router(server_router.router, tags=["Server"])

# NEW: Mount the WebSocket endpoint
app.websocket("/ws")(websocket_endpoint)