import type { FastifyInstance } from 'fastify';
import type { Sql } from '../db.js';
import type { Broker } from '@boocode/server/broker';

export function registerWebSocket(
  app: FastifyInstance,
  sql: Sql,
  broker: Broker,
): void {
  // Per-session streaming WebSocket. Clients connect here to receive live
  // inference frames (deltas, tool_calls, tool_results, message_complete).
  app.get<{ Params: { sessionId: string } }>(
    '/api/ws/sessions/:sessionId',
    { websocket: true },
    async (socket, req) => {
      const sessionId = req.params.sessionId;

      // Validate session exists
      const session = await sql<{ id: string }[]>`SELECT id FROM sessions WHERE id = ${sessionId}`;
      if (session.length === 0) {
        socket.send(JSON.stringify({ type: 'error', error: 'session not found' }));
        socket.close(1008, 'session not found');
        return;
      }

      // Send snapshot of existing messages so client can hydrate
      const messages = await sql<Record<string, unknown>[]>`
        SELECT id, session_id, chat_id, role, content, kind, tool_calls, tool_results, reasoning_parts, status, last_seq,
               tokens_used, ctx_used, ctx_max, started_at, finished_at, created_at, metadata,
               summary, tail_start_id, compacted_at
        FROM messages_with_parts
        WHERE session_id = ${sessionId}
        ORDER BY created_at ASC, id ASC
      `;
      socket.send(JSON.stringify({ type: 'snapshot', messages }));

      // Subscribe to broker for live frames
      const unsubscribe = broker.subscribe(sessionId, (frame) => {
        if (socket.readyState !== socket.OPEN) return;
        try {
          socket.send(JSON.stringify(frame));
        } catch (err) {
          app.log.warn({ err, sessionId }, 'ws send failed');
        }
      });

      socket.on('close', () => unsubscribe());
      socket.on('error', () => unsubscribe());
    },
  );
}