// Copyright Abridged, Inc. 2021,2024. All Rights Reserved.
// Node module: @collabland/aws
// This file is licensed under the MIT License.
// License text available at https://opensource.org/licenses/MIT

import {
  DeleteMessageBatchCommand,
  DeleteMessageBatchCommandInput,
  Message,
  ReceiveMessageCommand,
  ReceiveMessageCommandInput,
  SendMessageCommand,
  SendMessageCommandInput,
  SQSClient,
  SQSClientConfig,
} from '@aws-sdk/client-sqs';
import {
  fetchIteratorByBatch,
  fetchIteratorByPage,
  loggers,
  pMapByPage,
} from '@collabland/common';
import {
  getAWSAccountId,
  getAWSRegion,
  getAWSSQSRegion,
  getDefaultConfig,
} from '../common/types.js';

const {trace, error} = loggers('collabland:aws:sqs');

/**
 * Service client for AWS SQS
 */
export class SQSServiceClient {
  readonly sqs: SQSClient;

  constructor(options?: SQSClientConfig) {
    this.sqs = new SQSClient({
      ...getDefaultConfig(),
      region: getAWSSQSRegion(),
      ...options,
    });
  }

  /**
   * Send a message to the SQS queue
   * @param queueUrl - SQS queue url
   * @param message - Message
   * @param options - Send options
   * @returns
   */
  async send(
    queueUrl: string,
    message: string,
    options: Omit<SendMessageCommandInput, 'QueueUrl' | 'MessageBody'> = {},
  ) {
    const input: SendMessageCommandInput = {
      ...options,
      QueueUrl: queueUrl,
      MessageBody: message,
    };
    trace('Sending message to %s: %O', queueUrl, input);
    const cmd = new SendMessageCommand(input);
    const result = await this.sqs.send(cmd);
    return result;
  }

  async sendMessage(
    queueUrl: string,
    message: Message,
    options: Omit<
      SendMessageCommandInput,
      'QueueUrl' | 'MessageBody' | 'MessageAttributes'
    > = {},
  ) {
    const input: SendMessageCommandInput = {
      ...options,
      QueueUrl: queueUrl,
      MessageAttributes: message.MessageAttributes,
      MessageBody: message.Body,
    };
    trace('Sending message to %s: %O', queueUrl, input);
    const cmd = new SendMessageCommand(input);
    const result = await this.sqs.send(cmd);
    return result;
  }

  /**
   * Send a json object as message to SQS
   * @param queueUrl - SQS queue url
   * @param message - JSON message
   * @param options - Send options
   * @returns
   */
  async sendJson(
    queueUrl: string,
    message: object,
    options: Omit<SendMessageCommandInput, 'QueueUrl' | 'MessageBody'> = {},
  ) {
    return this.send(queueUrl, JSON.stringify(message), options);
  }

  /**
   * Receive and delete messages from the queue
   * @param queueUrl - SQS queue url
   * @param options - SQS receive options
   * @returns
   */
  async dequeueMessages(
    queueUrl: string,
    options: Omit<ReceiveMessageCommandInput, 'QueueUrl'> = {},
  ) {
    const output = await this.receive(queueUrl, options);
    await this.deleteMessages(queueUrl, output.Messages);
    return output;
  }

  /**
   * Delete messages from the queue
   * @param queueUrl - SQS queue url
   * @param messages - Messages to be deleted
   * @returns
   */
  async deleteMessages(queueUrl: string, messages?: Message[]) {
    if (messages == null || messages.length === 0) return [];
    return pMapByPage(
      messages,
      async page => {
        const deleteInput: DeleteMessageBatchCommandInput = {
          QueueUrl: queueUrl,
          Entries: page.map(m => ({
            Id: m.MessageId!,
            ReceiptHandle: m.ReceiptHandle,
          })),
        };
        const deleteOutput = await this.sqs.send(
          new DeleteMessageBatchCommand(deleteInput),
        );
        trace('DeleteMessageBatchCommandOutput %s: %O', queueUrl, deleteOutput);
        return deleteOutput;
      },
      // At most 10 messages in a batch
      {pageSize: 10},
    );
  }

  /**
   * Receive messages from the SQS queue
   * @param queueUrl - SQS queue url
   * @param options - Receive options
   * @returns
   */
  async receive(
    queueUrl: string,
    options: Omit<ReceiveMessageCommandInput, 'QueueUrl'>,
  ) {
    const input: ReceiveMessageCommandInput = {
      MessageAttributeNames: ['*'],
      ...options,
      QueueUrl: queueUrl,
    };
    trace('ReceiveMessageCommandInput %s: %O', queueUrl, input);
    const cmd = new ReceiveMessageCommand(input);
    const output = await this.sqs.send(cmd);
    trace('ReceiveMessageCommandOutput %s: %O', queueUrl, output);
    return output;
  }

  /**
   * Fetch messages from the queue as async iterator
   * @param queueUrl - Queue url
   * @param options - ReceiveMessage options
   * @returns
   */
  fetchMessages(
    queueUrl: string,
    options?: Omit<ReceiveMessageCommandInput, 'QueueUrl'> & {
      DeleteReceivedMessages?: boolean;
    },
  ) {
    return fetchIteratorByPage<Message, undefined>({
      next: async () => {
        try {
          const result = await this.receive(queueUrl, {
            WaitTimeSeconds: 20, // long polling
            MaxNumberOfMessages: 10,
            VisibilityTimeout: 300, // holding for 5 mins
            ...options,
          });
          const messages = result.Messages ?? [];
          if (options?.DeleteReceivedMessages && messages.length) {
            await this.deleteMessages(queueUrl, messages);
          }
          return {
            value: messages,
            done: false,
          };
        } catch (err) {
          error('Fail to receive/delete messages from %s: %O', queueUrl, err);
          return {value: [], done: false};
        }
      },
    });
  }

  fetchMessagesByBatch(
    queueUrl: string,
    options?: Omit<ReceiveMessageCommandInput, 'QueueUrl'> & {
      DeleteReceivedMessages?: boolean;
    },
  ) {
    const maxNumOfMessages = options?.MaxNumberOfMessages ?? 10;
    const maxNumOfMessagesPerTry = Math.min(10, maxNumOfMessages);
    const wait = options?.WaitTimeSeconds ?? maxNumOfMessages * 1;
    const maxTries = Math.ceil(maxNumOfMessages / 10);
    let waitPerTry = Math.min(20, Math.ceil(wait / maxTries)); // long polling
    // Make sure `WaitTimeSeconds` > connectionTimeout
    waitPerTry = Math.max(waitPerTry, 6);
    return fetchIteratorByBatch<Message, undefined>({
      next: async () => {
        try {
          let count = 0;
          let tries = 0;
          const messages: Message[] = [];
          while (count < maxNumOfMessages && tries < maxTries) {
            trace(
              'Receiving messages from %s (count: %s, tries: %s, maxTries: %s, maxNumOfMessages: %s)...',
              queueUrl,
              count,
              tries,
              maxTries,
              maxNumOfMessagesPerTry,
            );
            const result = await this.receive(queueUrl, {
              VisibilityTimeout: 300, // holding for 5 mins
              ...options,
              MaxNumberOfMessages: maxNumOfMessagesPerTry,
              WaitTimeSeconds: waitPerTry,
            });
            const received = result.Messages ?? [];
            messages.push(...received);
            count += received.length;
            tries++;
            if (options?.DeleteReceivedMessages && messages.length) {
              await this.deleteMessages(queueUrl, messages);
            }
          }
          trace(
            '%s messages received for the current batch from %s',
            messages.length,
            queueUrl,
          );
          return {
            value: messages,
            done: false,
          };
        } catch (err) {
          error('Fail to receive/delete messages from %s: %O', queueUrl, err);
          return {value: [], done: false};
        }
      },
    });
  }

  getMessageAttributes(message: Message, ...names: string[]) {
    const attributes: Record<string, string | number | Uint8Array | undefined> =
      {};
    if (names.length === 0) {
      names = Object.keys(message.MessageAttributes ?? {});
    }
    names.forEach(n => {
      const val = message.MessageAttributes?.[n];
      if (val != null) {
        if (val.DataType === 'String') {
          attributes[n] = val.StringValue;
        }
        if (val.DataType === 'Number') {
          attributes[n] = parseInt(val.StringValue!);
        }
        if (val.DataType === 'Binary') {
          attributes[n] = val.BinaryValue;
        }
      }
    });
    return attributes;
  }
}

/**
 * Get url for a SQS queue
 * @param queueName - Queue name
 * @param region - AWS region
 * @param accountId - AWS account id
 * @returns
 */
export function getSQSQueueUrl(
  queueName: string,
  region = getAWSRegion(),
  accountId = getAWSAccountId(),
) {
  const url = `https://sqs.${region}.amazonaws.com/${accountId}/${queueName}`;
  return url;
}