import {
  WebSocketGateway,
  WebSocketServer,
  SubscribeMessage,
  OnGatewayConnection,
  OnGatewayDisconnect,
  MessageBody,
  ConnectedSocket,
} from '@nestjs/websockets';
import { Server, Socket } from 'socket.io';
import { Logger } from '@nestjs/common';
import { AIVoiceService } from './ai-voice.service';
import { GeminiLiveProxyService } from './gemini-live-proxy.service';

@WebSocketGateway({
  namespace: '/ai-voice',
  cors: {
    origin: '*',
  },
})
export class AIVoiceGateway implements OnGatewayConnection, OnGatewayDisconnect {
  @WebSocketServer()
  server: Server;

  private readonly logger = new Logger(AIVoiceGateway.name);
  private readonly activeSessions = new Map<string, GeminiLiveProxyService>();
  private readonly sessionConversationIds = new Map<string, string | null>();

  constructor(private readonly aiVoiceService: AIVoiceService) {}

  async handleConnection(client: Socket) {
    const sessionId = client.id;
    this.logger.log(`AIVoiceGateway:handleConnection - client connected: ${sessionId}`);
  }

  async handleDisconnect(client: Socket) {
    const sessionId = client.id;
    this.logger.log(`AIVoiceGateway:handleDisconnect - client disconnected: ${sessionId}`);

    const proxyService = this.activeSessions.get(sessionId);
    if (proxyService) {
      await proxyService.disconnect();
      this.activeSessions.delete(sessionId);
      this.sessionConversationIds.delete(sessionId);
    }
  }

  @SubscribeMessage('connect-session')
  async handleConnectSession(
    @ConnectedSocket() client: Socket,
    @MessageBody() config: {
      modelName: string;
      voiceName: string;
      useSearch: boolean;
      conversationId?: string | null;
    },
  ) {
    const sessionId = client.id;
    const userId = client.handshake.query.userId as string;
    this.logger.log(`AIVoiceGateway:handleConnectSession - connecting session: ${sessionId}, userId: ${userId}`);

    try {
      if (this.activeSessions.has(sessionId)) {
        this.logger.warn(`AIVoiceGateway:handleConnectSession - session already exists: ${sessionId}`);
        await this.activeSessions.get(sessionId).disconnect();
      }

      const conversationId = config.conversationId || null;
      this.sessionConversationIds.set(sessionId, conversationId);

      const systemInstruction = await this.aiVoiceService.getSystemInstruction();
      this.logger.log(`AIVoiceGateway:handleConnectSession - using system instruction from backend, length: ${systemInstruction.length}`);

      const proxyService = new GeminiLiveProxyService(
        this.aiVoiceService.getApiKey(),
        (event, data) => {
          client.emit(event, data);
        },
        this.logger,
        async (role: 'user' | 'system', content: string) => {
          if (!userId) {
            this.logger.warn('AIVoiceGateway:handleConnectSession - userId not available, skipping message save');
            return null;
          }
          const currentConversationId = this.sessionConversationIds.get(sessionId) || null;
          const newConversationId = await this.aiVoiceService.addMessageToChatHistory(
            userId,
            currentConversationId,
            role,
            content,
          );
          if (newConversationId !== currentConversationId) {
            this.sessionConversationIds.set(sessionId, newConversationId);
          }
          return newConversationId;
        },
      );

      this.activeSessions.set(sessionId, proxyService);

      await proxyService.connect({
        ...config,
        systemInstruction,
      });

      client.emit('session-connected', { sessionId });
    } catch (error) {
      this.logger.error(`AIVoiceGateway:handleConnectSession - error: ${error.message}`, error.stack);
      client.emit('session-error', { error: error.message });
    }
  }

  @SubscribeMessage('disconnect-session')
  async handleDisconnectSession(@ConnectedSocket() client: Socket) {
    const sessionId = client.id;
    this.logger.log(`AIVoiceGateway:handleDisconnectSession - disconnecting session: ${sessionId}`);

    const proxyService = this.activeSessions.get(sessionId);
    if (proxyService) {
      await proxyService.disconnect();
      this.activeSessions.delete(sessionId);
      this.sessionConversationIds.delete(sessionId);
      client.emit('session-disconnected', { sessionId });
    }
  }

  @SubscribeMessage('send-audio')
  async handleSendAudio(
    @ConnectedSocket() client: Socket,
    @MessageBody() data: { media: { data: string; mimeType: string } },
  ) {
    const sessionId = client.id;
    const proxyService = this.activeSessions.get(sessionId);

    if (!proxyService) {
      this.logger.warn(`AIVoiceGateway:handleSendAudio - no active session: ${sessionId}`);
      return;
    }

    try {
      await proxyService.sendAudio(data.media);
    } catch (error) {
      this.logger.error(`AIVoiceGateway:handleSendAudio - error: ${error.message}`, error.stack);
      client.emit('session-error', { error: error.message });
    }
  }
}