import type { FastifyInstance, FastifyRequest } from 'fastify';
import type { WebSocket } from '@fastify/websocket';
import type { Authenticator } from '../auth/authenticator.js';
import { AgentHarness } from '../harness/agent-harness.js';
import type { InboundMessage } from '../types/messages.js';
import { randomUUID } from 'crypto';

import type { ProviderConfig } from '../llm/provider.js';

export interface WebSocketHandlerConfig {
  authenticator: Authenticator;
  providerConfig: ProviderConfig;
}

/**
 * WebSocket channel handler
 */
export class WebSocketHandler {
  private config: WebSocketHandlerConfig;
  private sessions = new Map<string, AgentHarness>();

  constructor(config: WebSocketHandlerConfig) {
    this.config = config;
  }

  /**
   * Register WebSocket routes
   */
  register(app: FastifyInstance): void {
    app.get(
      '/ws/chat',
      { websocket: true },
      async (socket: WebSocket, request: FastifyRequest) => {
        await this.handleConnection(socket, request, app);
      }
    );
  }

  /**
   * Handle WebSocket connection
   */
  private async handleConnection(
    socket: WebSocket,
    request: FastifyRequest,
    app: FastifyInstance
  ): Promise<void> {
    const logger = app.log;

    // Send initial connecting message
    socket.send(
      JSON.stringify({
        type: 'status',
        status: 'authenticating',
        message: 'Authenticating...',
      })
    );

    // Authenticate (this may take time if creating container)
    const authContext = await this.config.authenticator.authenticateWebSocket(request);
    if (!authContext) {
      logger.warn('WebSocket authentication failed');
      socket.send(
        JSON.stringify({
          type: 'error',
          message: 'Authentication failed',
        })
      );
      socket.close(1008, 'Authentication failed');
      return;
    }

    logger.info(
      { userId: authContext.userId, sessionId: authContext.sessionId },
      'WebSocket connection authenticated'
    );

    // Send workspace starting message
    socket.send(
      JSON.stringify({
        type: 'status',
        status: 'initializing',
        message: 'Starting your workspace...',
      })
    );

    // Create agent harness
    const harness = new AgentHarness({
      userId: authContext.userId,
      sessionId: authContext.sessionId,
      license: authContext.license,
      providerConfig: this.config.providerConfig,
      logger,
    });

    try {
      await harness.initialize();
      this.sessions.set(authContext.sessionId, harness);

      // Send connected message
      socket.send(
        JSON.stringify({
          type: 'connected',
          sessionId: authContext.sessionId,
          userId: authContext.userId,
          licenseType: authContext.license.licenseType,
          message: 'Connected to Dexorder AI',
        })
      );

      // Handle messages
      socket.on('message', async (data: Buffer) => {
        try {
          const payload = JSON.parse(data.toString());

          if (payload.type === 'message') {
            const inboundMessage: InboundMessage = {
              messageId: randomUUID(),
              userId: authContext.userId,
              sessionId: authContext.sessionId,
              content: payload.content,
              attachments: payload.attachments,
              timestamp: new Date(),
            };

            const response = await harness.handleMessage(inboundMessage);

            socket.send(
              JSON.stringify({
                type: 'message',
                ...response,
              })
            );
          }
        } catch (error) {
          logger.error({ error }, 'Error handling WebSocket message');
          socket.send(
            JSON.stringify({
              type: 'error',
              message: 'Failed to process message',
            })
          );
        }
      });

      // Handle disconnection
      socket.on('close', async () => {
        logger.info({ sessionId: authContext.sessionId }, 'WebSocket disconnected');
        await harness.cleanup();
        this.sessions.delete(authContext.sessionId);
      });

      socket.on('error', (error) => {
        logger.error({ error, sessionId: authContext.sessionId }, 'WebSocket error');
      });
    } catch (error) {
      logger.error({ error }, 'Failed to initialize agent harness');
      socket.close(1011, 'Internal server error');
      await harness.cleanup();
    }
  }
}
