Project_Velocity/backend/main.py

"""
Velocity — Unified FastAPI Backend
Covers: Catalyst (Meta Marketing), Sentinel (QD Engine), Vault (Trackable Links), Auth

GPU partitioning on AWS:
  - NemoClaw / Ollama → CUDA devices 0, 1  (enforced in nemoclaw.service systemd unit)
  - ComfyUI / Wan 2.2 → CUDA devices 2, 3  (enforced in comfyui.service systemd unit)
"""

import os
import json
import asyncio
import logging
from contextlib import asynccontextmanager
from datetime import UTC, datetime
from typing import Set

from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from dotenv import load_dotenv

from backend.api.routes_catalyst import router as catalyst_router
from backend.api.routes_crm import crm_router, analytics_router
from backend.api.routes_oracle import router as oracle_helper_router
from backend.auth.dependencies import (
    create_access_token, verify_password, get_current_user
)
from backend.db.pool import create_pool, close_pool
from backend.oracle.router_v1 import router as oracle_v1_router
from backend.routers.cctv import router as cctv_router
from backend.routers.scenes import router as scenes_router
from backend.routers.videos import router as videos_router
from backend.routers.vault import router as vault_router
from backend.routers.sentinel import router as sentinel_router, broadcast_sentinel_event

load_dotenv()

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("velocity.main")

# ── Lifespan: DB pool init / teardown ─────────────────────────────────────────

@asynccontextmanager
async def lifespan(app: FastAPI):
    # Startup
    try:
        app.state.db_pool = await create_pool()
        logger.info("asyncpg pool created")
    except Exception as exc:
        logger.error("Failed to create DB pool: %s", exc)
        app.state.db_pool = None

    app.state.broadcast_sentinel_event = broadcast_sentinel_event

    yield

    # Shutdown
    await close_pool()
    logger.info("asyncpg pool closed")

# ── App ───────────────────────────────────────────────────────────────────────

app = FastAPI(
    title="Velocity — Neural Core",
    description="Unified backend: Catalyst, Sentinel QD Engine, Vault, Oracle, Auth.",
    version="2.0.0",
    lifespan=lifespan,
)

# ── CORS ──────────────────────────────────────────────────────────────────────

origins = [o.strip() for o in os.getenv("CORS_ORIGINS", "http://localhost:5173").split(",")]

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

# ── Static asset serving (Vault files) ───────────────────────────────────────

ASSET_DIR = os.getenv("VELOCITY_ASSET_DIR", "/opt/dlami/nvme/assets")
if os.path.isdir(ASSET_DIR):
    app.mount("/assets", StaticFiles(directory=ASSET_DIR), name="assets")

# ── Routers ───────────────────────────────────────────────────────────────────

app.include_router(catalyst_router, prefix="/api/catalyst", tags=["Catalyst"])
app.include_router(crm_router, prefix="/api", tags=["CRM"])
app.include_router(analytics_router, prefix="/api/analytics", tags=["Analytics"])
app.include_router(oracle_helper_router, prefix="/api/oracle", tags=["Oracle"])
app.include_router(oracle_v1_router, prefix="/api/oracle/v1", tags=["Oracle V1"])
app.include_router(sentinel_router, prefix="/api/sentinel", tags=["Sentinel"])
app.include_router(cctv_router, prefix="/api/cctv", tags=["CCTV"])
app.include_router(scenes_router, prefix="/api/scenes", tags=["Scenes"])
app.include_router(videos_router, prefix="/api/videos", tags=["Videos"])
app.include_router(vault_router, prefix="/api/vault", tags=["Vault"])

# Public vault link (no /api prefix — shared externally with prospects)
from backend.routers.vault import router as public_vault_router
app.include_router(public_vault_router, prefix="/vault", tags=["Vault Public"])

# ── Auth endpoint ─────────────────────────────────────────────────────────────

from fastapi import HTTPException, status
from pydantic import BaseModel

class LoginRequest(BaseModel):
    email: str
    password: str

@app.post("/api/auth/login", tags=["Auth"])
async def login(body: LoginRequest):
    """
    Authenticate a user and return a JWT.
    Credentials are verified against the users_and_roles table.
    """
    from backend.db.pool import get_pool
    from fastapi import Request

    pool = app.state.db_pool
    if pool is None:
        raise HTTPException(status_code=503, detail="Database unavailable.")

    async with pool.acquire() as conn:
        row = await conn.fetchrow(
            "SELECT id::text, role, password_hash FROM users_and_roles WHERE email = $1 AND is_active = TRUE",
            body.email,
        )

    if not row or not verify_password(body.password, row["password_hash"]):
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Invalid email or password.",
        )

    token = create_access_token(user_id=row["id"], role=row["role"])
    return {"access_token": token, "token_type": "bearer", "expires_in": 28800}


@app.get("/api/auth/me", tags=["Auth"])
async def me(user=get_current_user):
    return {"user_id": user.user_id, "role": user.role}


# ── Catalyst WebSocket (preserved from v1) ────────────────────────────────────

class _CatalystManager:
    def __init__(self) -> None:
        self.active: Set[WebSocket] = set()

    async def connect(self, ws: WebSocket) -> None:
        await ws.accept()
        self.active.add(ws)

    def disconnect(self, ws: WebSocket) -> None:
        self.active.discard(ws)

    async def broadcast(self, payload: dict) -> None:
        dead: Set[WebSocket] = set()
        for ws in self.active:
            try:
                await ws.send_text(json.dumps(payload))
            except Exception:
                dead.add(ws)
        self.active -= dead


_catalyst_mgr = _CatalystManager()


class _CRMManager:
    def __init__(self) -> None:
        self.active: Set[WebSocket] = set()

    async def connect(self, ws: WebSocket) -> None:
        await ws.accept()
        self.active.add(ws)

    def disconnect(self, ws: WebSocket) -> None:
        self.active.discard(ws)

    async def broadcast(self, payload: dict) -> None:
        dead: Set[WebSocket] = set()
        for ws in self.active:
            try:
                await ws.send_text(json.dumps(payload))
            except Exception:
                dead.add(ws)
        self.active -= dead


_crm_mgr = _CRMManager()


@app.websocket("/ws/catalyst")
async def catalyst_ws(ws: WebSocket) -> None:
    await _catalyst_mgr.connect(ws)
    try:
        while True:
            data = await ws.receive_text()
            await ws.send_text(json.dumps({"type": "ack", "data": data}))
    except WebSocketDisconnect:
        _catalyst_mgr.disconnect(ws)


@app.websocket("/ws/crm")
async def crm_ws(ws: WebSocket) -> None:
    await _crm_mgr.connect(ws)
    await _crm_mgr.broadcast(
        {
                "type": "crm_presence",
                "connected_clients": len(_crm_mgr.active),
                "timestamp": datetime.now(UTC).isoformat(),
        }
    )
    try:
        while True:
            message = await ws.receive_text()
            await ws.send_text(json.dumps({"type": "crm_ack", "data": message}))
    except WebSocketDisconnect:
        _crm_mgr.disconnect(ws)


async def broadcast_live_event(event_type, message, campaign_name=None, value=None):
    payload = {
        "type": event_type,
        "message": message,
        "campaignName": campaign_name,
        "value": value,
        "timestamp": datetime.now(UTC).isoformat(),
    }
    await _catalyst_mgr.broadcast(payload)


app.state.broadcast_live_event = broadcast_live_event


async def broadcast_crm_event(payload: dict) -> None:
    enriched = {
        **payload,
        "timestamp": datetime.now(UTC).isoformat(),
    }
    await _crm_mgr.broadcast(enriched)


app.state.broadcast_crm_event = broadcast_crm_event


# ── Health ─────────────────────────────────────────────────────────────────────

@app.get("/health", tags=["Health"])
async def health() -> dict:
    pool = app.state.db_pool
    db_ok = pool is not None
    return {
        "status": "ok",
        "service": "velocity-backend",
        "version": "2.0.0",
        "db_pool": "connected" if db_ok else "unavailable",
        "timestamp": datetime.now(UTC).isoformat(),
    }