export type QueryResultDatasetRequest = {
  limit?: number;
  offset?: number;
  pageSize?: number;
  totalRows?: number;
};

export const QUERY_RESULT_DATASET_PAGE_SIZE = 1_000;

function maskSqlForDatasetClassifier(sql: string): string {
  let output = '';
  let index = 0;
  while (index < sql.length) {
    const char = sql[index];
    const next = sql[index + 1];
    if (char === '-' && next === '-') {
      const newline = sql.indexOf('\n', index + 2);
      output += ' ';
      index = newline === -1 ? sql.length : newline + 1;
      continue;
    }
    if (char === '/' && next === '*') {
      index += 2;
      let depth = 1;
      while (index < sql.length && depth > 0) {
        if (sql[index] === '/' && sql[index + 1] === '*') {
          depth += 1;
          index += 2;
          continue;
        }
        if (sql[index] === '*' && sql[index + 1] === '/') {
          depth -= 1;
          index += 2;
          continue;
        }
        index += 1;
      }
      output += ' ';
      continue;
    }
    if (char === "'" || char === '"') {
      const quote = char;
      output += ' ';
      index += 1;
      while (index < sql.length) {
        const quoted = sql[index];
        if (quoted === quote && sql[index + 1] === quote) {
          index += 2;
          continue;
        }
        index += 1;
        if (quoted === quote) break;
      }
      continue;
    }
    output += char;
    index += 1;
  }
  return output;
}

export function isCustomerDbDatasetTool(toolId: string): boolean {
  return /^(?:customer_db_)?query_customer_db$/.test(toolId);
}

export function isSnowflakeQueryDatasetTool(toolId: string): boolean {
  return /^snowflake_(?:snowflake_)?run_(?:semantic_)?query$/.test(toolId);
}

function isSnowflakeSemanticQueryDatasetTool(toolId: string): boolean {
  return /^snowflake_(?:snowflake_)?run_semantic_query$/.test(toolId);
}

export function isQueryResultDatasetTool(toolId: string): boolean {
  return isCustomerDbDatasetTool(toolId) || isSnowflakeQueryDatasetTool(toolId);
}

export function isQueryResultDatasetOperation(input: {
  provider?: string | null;
  operation?: string | null;
}): boolean {
  return (
    (input.provider === 'customer_db' &&
      input.operation === 'query_customer_db') ||
    (input.provider === 'snowflake' &&
      (input.operation === 'snowflake_run_query' ||
        input.operation === 'snowflake_run_semantic_query'))
  );
}

function isReadOnlySelectOrWith(sql: string): boolean {
  const normalized = maskSqlForDatasetClassifier(sql).trimStart().toLowerCase();
  return (
    (normalized.startsWith('select') &&
      !/\bselect\b[\s\S]*?\binto\b/i.test(normalized)) ||
    (normalized.startsWith('with') &&
      !/\b(insert\s+into|update|delete\s+from|merge\s+into|select\b[\s\S]*?\binto|create|alter|drop|truncate)\b/i.test(
        normalized,
      ))
  );
}

export function isQueryResultDatasetReadRequest(
  toolId: string,
  input: Record<string, unknown>,
): boolean {
  if (!isQueryResultDatasetTool(toolId)) return false;
  if (isSnowflakeQueryDatasetTool(toolId) && input.write === true) return false;
  if (isSnowflakeSemanticQueryDatasetTool(toolId)) return true;
  const sql =
    typeof input.sql === 'string'
      ? input.sql
      : typeof input.query === 'string'
        ? input.query
        : '';
  return isReadOnlySelectOrWith(sql);
}

export const isCustomerDbDatasetReadRequest =
  isQueryResultDatasetReadRequest;