import type { FastifyInstance } from 'fastify';
import type { Sql } from '../db.js';
import type { Broker } from '../services/broker.js';
import type { Message } from '../types/api.js';
import { MESSAGE_COLUMNS } from '../services/message-columns.js';
import { loadAgentSnapshot } from '../services/session-snapshots.js';

export function registerWebSocket(
  app: FastifyInstance,
  sql: Sql,
  broker: Broker
): void {
  app.get<{ Params: { id: string } }>(
    '/api/ws/sessions/:id',
    { websocket: true },
    async (socket, req) => {
      const sessionId = req.params.id;

      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;
      }

      // v1.11: snapshot includes compaction fields so MessageBubble can
      // render the SummaryCard for summary=true rows on first connect.
      // v1.13.1-B: reads tool_calls/tool_results via the parts-merged view.
      const messages = await sql<Message[]>`
        SELECT ${sql.unsafe(MESSAGE_COLUMNS)}
        FROM messages_with_parts
        WHERE session_id = ${sessionId}
        ORDER BY created_at ASC, id ASC
      `;
      socket.send(JSON.stringify({ type: 'snapshot', messages }));

      // v2.7.x: on reconnect, restore agent snapshot state so the frontend
      // knows there's an ongoing agent turn. Best-effort per chat; most
      // sessions won't have any snapshots.
      const chats = await sql<{ id: string }[]>`SELECT id FROM chats WHERE session_id = ${sessionId}`;
      for (const chat of chats) {
        const agentSnapshot = await loadAgentSnapshot(sql, chat.id).catch(() => null);
        if (agentSnapshot) {
          socket.send(JSON.stringify({
            type: 'agent_snapshot',
            chat_id: chat.id,
            agent: agentSnapshot.agent,
            model: agentSnapshot.model,
            mode: agentSnapshot.mode,
            turn_number: agentSnapshot.turn_number,
          }));
        }
      }

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

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