import { AttributeValue, DynamoDBClient, QueryCommand, QueryCommandOutput } from '@aws-sdk/client-dynamodb';
import { JobExecutionStatus } from '../../constants/JobExecutionStatus.constants';
import { JobExecutionDto, JobExecutionsByTenant } from '../../model/JobExecution';
import { parsePartitionKey } from '../../ParsePartitionKey';
import { convertDynamoDbItemToJobExecution } from '../JobExecutionRepository/utils';
import { ENTITY_TYPES } from '../../constants/Repository.constants';

export class MultiTenantJobExecutionRepository {
  public constructor(
    private readonly dynamoDbClient: DynamoDBClient,
    private readonly tableName: string,
  ) {}

  public async findByStatus(status: JobExecutionStatus): Promise<JobExecutionsByTenant> {
    const result: JobExecutionsByTenant = {};
    let startKey: Record<string, AttributeValue> | undefined = undefined;

    do {
      const res: QueryCommandOutput = await this.dynamoDbClient.send(
        new QueryCommand({
          KeyConditionExpression: 'gsi1pk = :gsi1pk',
          ExpressionAttributeValues: {
            ':gsi1pk': { S: `${ENTITY_TYPES.jobExecution}~${status}` },
          },
          TableName: this.tableName,
          IndexName: 'gsi1',
          ExclusiveStartKey: startKey,
        }),
      );

      res.Items?.forEach((item) => {
        const { tenant } = parsePartitionKey(item.pk.S || '');
        result[tenant] ||= [];
        result[tenant].push(convertDynamoDbItemToJobExecution(item as JobExecutionDto));
      });

      startKey = res.LastEvaluatedKey;
    } while (startKey);

    return result;
  }
}