# backend/main.py
import json
import asyncio
from contextlib import asynccontextmanager
from fastapi import FastAPI, WebSocket, Depends
from sqlalchemy import text
from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
import psycopg

# Global store: tenant_id -> set of WebSocket connections
tenant_connections: dict[str, set[WebSocket]] = {}

async def get_db(request):
    async with AsyncSession(engine) as session:
        yield session

# FastAPI lifespan: start listener in background
@asynccontextmanager
async def lifespan(app: FastAPI):
    task = asyncio.create_task(listen_for_notifications())
    yield
    task.cancel()

app = FastAPI(lifespan=lifespan)

async def listen_for_notifications():
    """
    Connect to PostgreSQL directly and listen for NOTIFY events.
    This runs in the background for the entire app lifetime.
    """
    conn = await psycopg.AsyncConnection.connect(DATABASE_URL)

    # Listen on all tenant channels dynamically
    # In production, you'd track which channels are actually subscribed
    async with conn.cursor() as cur:
        # We'll use LISTEN with a pattern-like approach
        await cur.execute("LISTEN notifications")

    try:
        async for notify in conn.notifies():
            payload = json.loads(notify.payload)
            tenant_id = payload.get("tenant_id")

            if tenant_id in tenant_connections:
                # Broadcast to all WebSockets for this tenant
                disconnected = set()
                for websocket in tenant_connections[tenant_id]:
                    try:
                        await websocket.send_json({
                            "type": payload.get("type"),
                            "data": payload.get("data"),
                        })
                    except Exception:
                        disconnected.add(websocket)

                # Clean up dead connections
                tenant_connections[tenant_id] -= disconnected
                if not tenant_connections[tenant_id]:
                    del tenant_connections[tenant_id]
    finally:
        await conn.close()

@app.websocket("/ws/notifications")
async def websocket_notifications(websocket: WebSocket):
    await websocket.accept()

    # Extract tenant from JWT or query params
    tenant_id = websocket.query_params.get("tenant_id")

    if not tenant_id:
        await websocket.close(code=1008, reason="Missing tenant_id")
        return

    # Register this connection
    if tenant_id not in tenant_connections:
        tenant_connections[tenant_id] = set()
    tenant_connections[tenant_id].add(websocket)

    try:
        # Keep connection open, listen for client disconnects
        while True:
            await websocket.receive_text()
    except Exception:
        pass
    finally:
        tenant_connections[tenant_id].discard(websocket)

@app.post("/api/notifications")
async def send_notification(
    tenant_id: str,
    message: str,
    db: AsyncSession = Depends(get_db),
):
    """
    Example endpoint that triggers a notification.
    The NOTIFY happens inside the transaction.
    """
    # Create the notification record
    await db.execute(
        text("""
            INSERT INTO notifications (tenant_id, message, created_at)
            VALUES (:tenant_id, :message, NOW())
        """),
        {"tenant_id": tenant_id, "message": message},
    )

    # Send the PostgreSQL notification
    payload = json.dumps({
        "tenant_id": tenant_id,
        "type": "notification",
        "data": {"message": message},
    })
    await db.execute(
        text("SELECT pg_notify('notifications', :payload)"),
        {"payload": payload},
    )

    await db.commit()

    return {"ok": True}

// src/hooks/useNotifications.ts
import { useEffect, useRef } from "react";

export function useNotifications(tenantId: string, onMessage: (data: any) => void) {
  const wsRef = useRef<WebSocket | null>(null);

  useEffect(() => {
    const protocol = window.location.protocol === "https:" ? "wss:" : "ws:";
    const wsUrl = `${protocol}//${window.location.host}/ws/notifications?tenant_id=${tenantId}`;

    const ws = new WebSocket(wsUrl);

    ws.onopen = () => {
      console.log("✓ Real-time connection established");
    };

    ws.onmessage = (event) => {
      const message = JSON.parse(event.data);
      onMessage(message);
    };

    ws.onerror = () => {
      console.error("WebSocket error");
    };

    wsRef.current = ws;

    return () => {
      ws.close();
    };
  }, [tenantId, onMessage]);
}

// Usage in a component
export function NotificationCenter({ tenantId }: { tenantId: string }) {
  const [notifications, setNotifications] = useState<any[]>([]);

  useNotifications(tenantId, (message) => {
    setNotifications((prev) => [message, ...prev]);
  });

  return (
    <div className="space-y-2">
      {notifications.map((notif, i) => (
        <div key={i} className="p-3 bg-blue-50 rounded border-l-4 border-blue-400">
          {notif.data.message}
        </div>
      ))}
    </div>
  );
}

# backend/celery_task.py (or similar background job)
async def process_ai_feature(tenant_id: str, job_id: str, db: AsyncSession):
    # ... do AI work with Claude ...
    result = await call_claude_api(...)

    # Save result
    await db.execute(
        text("""
            UPDATE ai_jobs 
            SET result = :result, status = 'completed'
            WHERE id = :job_id AND tenant_id = :tenant_id
        """),
        {"result": json.dumps(result), "job_id": job_id, "tenant_id": tenant_id},
    )

    # Notify the tenant in real-time
    payload = json.dumps({
        "tenant_id": tenant_id,
        "type": "ai_job_complete",
        "data": {"job_id": job_id, "result": result},
    })
    await db.execute(
        text("SELECT pg_notify('notifications', :payload)"),
        {"payload": payload},
    )

    await db.commit()

python
# Wrong: connection pool conflict
async with engine.begin() as conn:
    await conn.execute(text("LISTEN notifications"))