import { Inject } from '@nestjs/common';
import {
  WebSocketGateway,
  WebSocketServer,
  SubscribeMessage,
  OnGatewayConnection,
  OnGatewayDisconnect,
  ConnectedSocket,
  MessageBody,
} from '@nestjs/websockets';
import { Server, Socket } from 'socket.io';
import {
  TenantContextService,
  TENANT_FINDER,
  DEFAULT_TENANT_ID,
} from '@iconsulting/shared';
import type { ITenantFinder } from '@iconsulting/shared';
import { ConversationService, FileAttachment } from '../../application/services/conversation.service';

interface SendMessagePayload {
  conversationId: string;
  content: string;
  attachments?: FileAttachment[];
}

interface ConnectionInfo {
  userId: string;
  tenantId: string;
}

@WebSocketGateway({
  cors: {
    origin: process.env.CORS_ORIGINS?.split(',') || ['http://localhost:5173'],
    credentials: true,
  },
  namespace: '/ws/conversation',
  path: '/ws/conversation/socket.io',
})
export class ConversationGateway
  implements OnGatewayConnection, OnGatewayDisconnect
{
  @WebSocketServer()
  server: Server;

  // Map socket ID to connection info (userId + tenantId)
  private connections = new Map<string, ConnectionInfo>();

  constructor(
    private conversationService: ConversationService,
    private readonly tenantContext: TenantContextService,
    @Inject(TENANT_FINDER) private readonly tenantFinder: ITenantFinder,
  ) {}

  async handleConnection(client: Socket) {
    // Extract user ID from query or headers
    const userId =
      (client.handshake.query.userId as string) ||
      (client.handshake.headers['x-user-id'] as string);

    if (!userId) {
      console.log(`Client ${client.id} connected without user ID`);
      client.emit('error', { message: 'User ID is required' });
      client.disconnect();
      return;
    }

    // Extract tenant ID from query or headers, fall back to default
    const tenantId =
      (client.handshake.query.tenantId as string) ||
      (client.handshake.headers['x-tenant-id'] as string) ||
      DEFAULT_TENANT_ID;

    this.connections.set(client.id, { userId, tenantId });
    console.log(`Client ${client.id} connected as user ${userId} (tenant: ${tenantId})`);

    client.emit('connected', { userId, socketId: client.id });
  }

  handleDisconnect(client: Socket) {
    const conn = this.connections.get(client.id);
    this.connections.delete(client.id);
    console.log(`Client ${client.id} (user ${conn?.userId}) disconnected`);
  }

  @SubscribeMessage('message')
  async handleMessage(
    @ConnectedSocket() client: Socket,
    @MessageBody() payload: SendMessagePayload,
  ) {
    const conn = this.connections.get(client.id);

    if (!conn) {
      client.emit('error', { message: 'Not authenticated' });
      return;
    }

    const { conversationId, content, attachments } = payload;

    if (!conversationId || !content) {
      client.emit('error', { message: 'conversationId and content are required' });
      return;
    }

    // Resolve tenant context for this connection
    const tenant = await this.tenantFinder.findById(conn.tenantId);
    if (!tenant) {
      client.emit('error', { message: 'Invalid tenant', conversationId });
      return;
    }

    // Run the message handler inside tenant context (AsyncLocalStorage)
    await this.tenantContext.runAsync(tenant, async () => {
      this.tenantContext.setCurrentUserId(conn.userId);

      try {
        // Generate unique message ID for this response
        const messageId = `msg_${Date.now()}_${Math.random().toString(36).substr(2, 9)}`;

        // Emit stream start
        client.emit('stream_start', { messageId, conversationId });

        let chunkIndex = 0;

        // Stream the response
        for await (const chunk of this.conversationService.sendMessage({
          conversationId,
          userId: conn.userId,
          content,
          attachments,
        })) {
          if (chunk.type === 'text' && chunk.content) {
            client.emit('stream_chunk', {
              messageId,
              conversationId,
              content: chunk.content,
              index: chunkIndex++,
            });
          } else if (chunk.type === 'tool_use') {
            client.emit('tool_call', {
              messageId,
              conversationId,
              tool: chunk.toolName,
              input: chunk.toolInput,
            });
          } else if (chunk.type === 'tool_result') {
            client.emit('tool_result', {
              messageId,
              conversationId,
              tool: chunk.toolName,
              result: chunk.toolResult,
            });
          } else if (chunk.type === 'agent_start') {
            client.emit('agent_start', {
              messageId,
              conversationId,
              agentType: chunk.agentType,
              agentName: chunk.agentName,
              description: chunk.description,
            });
          } else if (chunk.type === 'agent_complete') {
            client.emit('agent_complete', {
              messageId,
              conversationId,
              agentType: chunk.agentType,
              agentName: chunk.agentName,
              durationMs: chunk.durationMs,
              success: chunk.success,
            });
          } else if (chunk.type === 'coordinator_thinking') {
            client.emit('coordinator_thinking', {
              messageId,
              conversationId,
              phase: chunk.phase,
              message: chunk.message,
            });
          } else if (chunk.type === 'end') {
            client.emit('stream_end', {
              messageId,
              conversationId,
              isComplete: true,
              inputTokens: chunk.inputTokens,
              outputTokens: chunk.outputTokens,
            });
          }
        }
      } catch (error) {
        console.error('Error processing message:', error);
        client.emit('error', {
          message: error instanceof Error ? error.message : 'Failed to process message',
          conversationId,
        });
      }
    }, conn.userId);
  }

  @SubscribeMessage('typing_start')
  handleTypingStart(
    @ConnectedSocket() client: Socket,
    @MessageBody() payload: { conversationId: string },
  ) {
    // Could be used to show typing indicator to admin monitoring
    console.log(`User typing in conversation ${payload.conversationId}`);
  }

  @SubscribeMessage('typing_end')
  handleTypingEnd(
    @ConnectedSocket() client: Socket,
    @MessageBody() payload: { conversationId: string },
  ) {
    console.log(`User stopped typing in conversation ${payload.conversationId}`);
  }
}
