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;

      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, model, 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());
    },
  );

  // User-channel WS: run-level orchestrator frames (flow_run_started,
  // flow_run_step_updated) published by the flow-runner via
  // broker.publishUserFrame('default'). Mirrors the BooChat server pattern.
  app.get('/api/ws/user', { websocket: true }, async (socket) => {
    const user = 'default';
    const unsubscribe = broker.subscribeUser(user, (frame) => {
      if (socket.readyState !== socket.OPEN) return;
      try {
        socket.send(JSON.stringify(frame));
      } catch (err) {
        app.log.warn({ err, user }, 'user ws send failed');
      }
    });
    socket.on('close', () => unsubscribe());
    socket.on('error', () => unsubscribe());
  });
}
