import { DynamoDBDocument, TransactWriteCommandInput } from '@aws-sdk/lib-dynamodb';
import { randomUUID } from 'crypto';
import { TransactionError } from '../error/TransactionError';
import { trace } from '@opentelemetry/api';

const tracer = trace.getTracer('db-lib:DynamoDbManager');

export type Transaction = {
  id?: string;
};

type TransactionItems = TransactWriteCommandInput['TransactItems'];
export type TransactionItem = NonNullable<TransactionItems>[number];

export class DynamoDbManager<TRepositories> {
  private transactionItems: Record<string, TransactionItem[]>;
  public repositories: TRepositories;

  constructor(
    public client: DynamoDBDocument,
    repositoryCreator: (dbManager: DynamoDbManager<TRepositories>) => TRepositories,
  ) {
    this.transactionItems = {};
    this.repositories = repositoryCreator(this);
  }

  public async executeInTransaction<T>(func: (transaction: Transaction) => Promise<T>): Promise<T> {
    const transactionId = randomUUID();

    try {
      this.startTransaction(transactionId);
      const value = await func({ id: transactionId });
      await this.executeTransaction(transactionId);
      return value;
    } finally {
      this.closeTransaction(transactionId);
    }
  }

  public addWriteTransactionItem(transactionId: string, item: TransactionItem) {
    if (this.transactionItems[transactionId] === undefined) {
      throw new TransactionError(`No items in transaction '${transactionId}' to add transaction item to`);
    }

    this.transactionItems[transactionId].push(item);
  }

  private async executeTransaction(transactionId: string) {
    if (this.transactionItems[transactionId] === undefined) {
      throw new TransactionError(`No items in transaction '${transactionId}' to execute`);
    }
    if (!this.transactionItems[transactionId].length) {
      return;
    }

    return tracer.startActiveSpan(
      'executeTransaction',
      {
        attributes: {
          'transactionItems.length': this.transactionItems[transactionId].length,
        },
      },
      async (span) => {
        try {
          return await this.client.transactWrite({
            ClientRequestToken: transactionId,
            TransactItems: this.transactionItems[transactionId],
          });
        } finally {
          span.end();
        }
      },
    );
  }

  private startTransaction(transactionId: string) {
    if (this.transactionItems[transactionId] !== undefined) {
      throw new TransactionError(`Transaction '${transactionId}' already started`);
    }
    this.transactionItems[transactionId] = [];
  }

  private closeTransaction(transactionId: string) {
    delete this.transactionItems[transactionId];
  }
}