import {
  WebSocketGateway,
  WebSocketServer,
  OnGatewayConnection,
  OnGatewayDisconnect,
  SubscribeMessage,
  MessageBody,
  ConnectedSocket,
} from '@nestjs/websockets';
import { Logger } from '@nestjs/common';
import { Server } from 'ws';
import type { WebSocket } from 'ws';

/**
 * Raw WebSocket gateway for agent stream events.
 *
 * Uses @nestjs/platform-ws (WsAdapter) instead of Socket.IO so that
 * plain WebSocket clients (Flutter, curl, etc.) can connect directly
 * without the Socket.IO handshake protocol.
 *
 * Path: ws://<host>:<port>/ws/agent
 */
@WebSocketGateway({ path: '/ws/agent' })
export class AgentStreamGateway implements OnGatewayConnection, OnGatewayDisconnect {
  private readonly logger = new Logger(AgentStreamGateway.name);

  @WebSocketServer()
  server!: Server;

  /** sessionId → Set of subscribed WebSocket clients */
  private readonly sessionClients = new Map<string, Set<WebSocket>>();
  /** client → Set of sessionIds it subscribed to (for cleanup on disconnect) */
  private readonly clientSessions = new Map<WebSocket, Set<string>>();
  /** sessionId → buffered events (replayed on late subscription) */
  private readonly sessionBuffer = new Map<string, any[]>();
  /** Max events to buffer per session before discarding oldest */
  private readonly MAX_BUFFER_SIZE = 500;
  /** TTL for orphan buffers (ms) — cleaned up if no subscriber after this time */
  private readonly BUFFER_TTL_MS = 120_000;

  handleConnection(client: WebSocket) {
    this.logger.log('Client connected');
    this.clientSessions.set(client, new Set());

    // Handle raw messages (WsAdapter delivers parsed JSON to @SubscribeMessage
    // only when the payload has an "event" field; otherwise we handle manually)
    client.on('message', (raw: Buffer | string) => {
      try {
        const msg = JSON.parse(raw.toString());
        if (msg.event === 'subscribe_session') {
          this.handleSubscribeSession(client, msg.data ?? msg);
        }
      } catch {
        // Ignore non-JSON or malformed messages
      }
    });
  }

  handleDisconnect(client: WebSocket) {
    this.logger.log('Client disconnected');
    const sessions = this.clientSessions.get(client);
    if (sessions) {
      for (const sid of sessions) {
        const clients = this.sessionClients.get(sid);
        if (clients) {
          clients.delete(client);
          if (clients.size === 0) {
            this.sessionClients.delete(sid);
          }
        }
      }
    }
    this.clientSessions.delete(client);
  }

  handleSubscribeSession(client: WebSocket, payload: { sessionId: string; taskId?: string }) {
    const { sessionId } = payload;
    if (!sessionId) return;

    // Add client to session room
    if (!this.sessionClients.has(sessionId)) {
      this.sessionClients.set(sessionId, new Set());
    }
    this.sessionClients.get(sessionId)!.add(client);
    this.clientSessions.get(client)?.add(sessionId);

    // Acknowledge subscription
    this.sendTo(client, {
      event: 'subscribed',
      data: { sessionId, taskId: payload.taskId },
    });

    // Replay any buffered events that arrived before this client subscribed
    const buffered = this.sessionBuffer.get(sessionId);
    if (buffered && buffered.length > 0) {
      this.logger.log(`Replaying ${buffered.length} buffered events for session ${sessionId}`);
      for (const evt of buffered) {
        this.sendTo(client, { event: 'stream_event', data: evt });
      }
      // Clear buffer after replay (all events delivered)
      this.sessionBuffer.delete(sessionId);
    }

    this.logger.log(`Client subscribed to session ${sessionId}`);
  }

  /**
   * Emits a stream event to all WebSocket clients subscribed to the session.
   * If no clients are subscribed yet, the event is buffered for late subscribers.
   */
  emitStreamEvent(sessionId: string, event: any) {
    const clients = this.sessionClients.get(sessionId);
    const hasSubscribers = clients && clients.size > 0;

    if (hasSubscribers) {
      const message = JSON.stringify({ event: 'stream_event', data: event });
      for (const client of clients) {
        if (client.readyState === 1 /* WebSocket.OPEN */) {
          client.send(message);
        }
      }
    } else {
      // No subscribers yet — buffer the event for late joiners
      if (!this.sessionBuffer.has(sessionId)) {
        this.sessionBuffer.set(sessionId, []);
        // Schedule cleanup in case no one ever subscribes
        setTimeout(() => {
          if (this.sessionBuffer.has(sessionId)) {
            this.sessionBuffer.delete(sessionId);
          }
        }, this.BUFFER_TTL_MS);
      }
      const buffer = this.sessionBuffer.get(sessionId)!;
      buffer.push(event);
      if (buffer.length > this.MAX_BUFFER_SIZE) {
        buffer.shift(); // Drop oldest event
      }
    }

    // Clean up buffer on terminal events
    if (event.type === 'completed' || event.type === 'error') {
      // Give late subscribers a small window to connect and receive replay
      setTimeout(() => {
        this.sessionBuffer.delete(sessionId);
      }, 10_000);
    }
  }

  private sendTo(client: WebSocket, data: any) {
    if (client.readyState === 1) {
      client.send(JSON.stringify(data));
    }
  }
}
