// kqlClient.ts

import { Client as KustoClient, KustoConnectionStringBuilder, ClientRequestProperties } from 'azure-kusto-data';
import { Logger } from 'winston';
import { v4 as uuidv4 } from 'uuid';
import { DefaultAzureCredential } from '@azure/identity';
import { Sema } from 'async-sema';

export class KqlClient {
  private client: KustoClient;
  private databaseName: string;
  private logger: Logger;
  private semaphore: Sema = new Sema(1);

  constructor(connectionString: string, databaseName: string, logger: Logger) {
    const kcsb = KustoConnectionStringBuilder.withTokenCredential(
      connectionString,
      new DefaultAzureCredential(),
    );
    this.client = new KustoClient(kcsb);
    this.databaseName = databaseName;
    this.logger = logger;
  }

  public async queryInChunks(
    tableName: string,
    filters: Record<string, string | Date | [Date, Date]>,
    onRows: (rows: Record<string, unknown>[]) => void,
    chunkSize : number = 10000
  ): Promise<void> {
    let hasMoreData = true;
    let lastTime: Date | null = null;
    let lastId: string | null = null;

    await this.semaphore.acquire();
    try
    {
      while (hasMoreData) {
        // Build the query with filters and pagination
        let query = `${tableName} | extend ___tid=strcat(tostring(___time), ___id) `;
        const filterStrings = [];

        // Apply provided filters
        for (const [key, value] of Object.entries(filters)) {
          if (key === '___time') {
            if (Array.isArray(value)) {
              const [fromTime, toTime] = value;
              filterStrings.push(
                `strcmp( ___tid, "${fromTime.toISOString()}") >= 0 and strcmp(___tid, "${toTime.toISOString()}") <= 0`
              );
            } else {
              filterStrings.push(`strcmp(___tid,strcat(datetime("${value.toISOString()}"), ___id)) >= 0`);
            }
          } else if (typeof value === 'string') {
            filterStrings.push(`${key} startswith "${value}"`);
          } else {
            filterStrings.push(`${key} == ${value}`);
          }
        }

        // Add pagination filters
        if (lastTime) {
          const timeFilter = `strcmp(___tid,strcat(datetime("${lastTime.toISOString()}"),___id)) > 0`;
          filterStrings.push(timeFilter);
        }

        if (filterStrings.length > 0) {
          query += ` | where ${filterStrings.join(' and ')}`;
        }

        query += ` | top ${chunkSize} by ___tid asc | project-away ___tid`;

        this.logger.info(`Executing KQL Query: ${query}`);

        try {
          const queryResults = await this.client.execute(
            this.databaseName,
            query,
            new ClientRequestProperties()
          );
          const table = queryResults.primaryResults[0];
          const rows = table.toJSON();

          onRows(rows.data.map((row: Record<string, unknown>) => row));

          if (rows.data.length < chunkSize) {
            hasMoreData = false;
            break;
          }

          // Update lastTime and lastId for the next chunk
          const lastRow = rows.data[rows.data.length - 1];
          lastTime = new Date(String(lastRow['___time']));
          lastId = String(lastRow['___id']);
        } catch (error) {
          this.logger.error(`Error executing query: ${error}`);
          throw error;
        }
      }
    } finally {
      this.semaphore.release();
    }
  }
}