import { Inject, Provide, Scope, ScopeEnum } from '@midwayjs/core';
import { InjectEntityModel } from '@midwayjs/typeorm';
import { createHash } from 'crypto';
import { Repository } from 'typeorm';
import { NetaClawDataSourceEntity } from '../entity/data_source.js';
import { NetaClawDataSourceQueryAuditEntity } from '../entity/data_source_query_audit.js';
import { MysqlPoolManager } from './mysql_pool.js';

export interface MysqlSqlGuardOptions {
  allowedTables?: string[];
  blockedTables?: string[];
  maxJoinTables?: number;
  maxRows?: number;
  maskedColumns?: Record<string, 'hash' | 'partial' | 'redact'>;
}

export interface MysqlSqlGuardResult {
  ok: boolean;
  reason?: string;
  tables?: string[];
  limitedSql?: string;
  appendedLimit?: boolean;
}

export interface ExecuteReadOnlyParams {
  source: NetaClawDataSourceEntity;
  sql: string;
  agentId?: number | null;
  userId?: number | null;
  toolCallId?: string | null;
}

export interface MysqlReadOnlyQueryResult {
  columns: string[];
  rows: Record<string, unknown>[];
  rowCount: number;
  truncated: boolean;
  elapsedMs: number;
  sql: string;
}

const METADATA_SQL_RE = /^(show|describe|desc|explain)\b/i;
const DML_SQL_RE = /^(insert|update|delete|replace|merge|truncate)\b/i;
const DDL_SQL_RE = /^(create|alter|drop|rename)\b/i;
const DCL_SQL_RE = /^(grant|revoke)\b/i;
const TRANSACTION_SQL_RE = /^(begin|start\s+transaction|commit|rollback|savepoint|release\s+savepoint|set\s+transaction)\b/i;
const COMMENT_RE = /(--|#|\/\*|\*\/)/;
const DANGEROUS_FUNCTION_RE = /\b(load_file|sleep|benchmark|get_lock|release_lock|is_free_lock|is_used_lock)\s*\(/i;
const LOCKING_READ_RE = /\b(for\s+(update|share)|lock\s+in\s+share\s+mode)\b/i;
const CLAUSE_BOUNDARY_RE = /\b(where|group\s+by|order\s+by|having|limit|procedure|into)\b/i;
const TABLE_REF_RE =
  /\b(?:from|join)\s+((?:`[^`]+`|[A-Za-z_][\w$]*)(?:\s*\.\s*(?:`[^`]+`|[A-Za-z_][\w$]*))?)/gi;
const RESERVED_ALIAS_WORDS = new Set([
  'where',
  'join',
  'inner',
  'left',
  'right',
  'full',
  'cross',
  'on',
  'using',
  'group',
  'order',
  'having',
  'limit',
]);

interface MysqlTableReference {
  table: string;
  schemaQualified: boolean;
  alias?: string;
}

function clampNumber(value: number | undefined, fallback: number, min: number, max: number) {
  if (!Number.isFinite(value)) {
    return fallback;
  }
  return Math.min(max, Math.max(min, Math.trunc(value as number)));
}

function elapsedSince(startedAt: bigint): number {
  return Number((process.hrtime.bigint() - startedAt) / BigInt(1_000_000));
}

function normalizeTableName(rawName: string) {
  const parts = rawName
    .split('.')
    .map(part => part.trim().replace(/^`|`$/g, ''))
    .filter(Boolean);
  return (parts[parts.length - 1] ?? '').toLowerCase();
}

function normalizeIdentifier(rawName: string) {
  return rawName.trim().replace(/^`|`$/g, '').toLowerCase();
}

function unique(values: string[]) {
  return [...new Set(values)];
}

function deny(reason: string, tables?: string[]): MysqlSqlGuardResult {
  return { ok: false, reason, ...(tables ? { tables } : {}) };
}

export function extractMysqlTables(sql: string): string[] {
  return unique(extractMysqlTableReferences(sql).map(ref => ref.table));
}

function extractMysqlTableReferences(sql: string): MysqlTableReference[] {
  const tables: string[] = [];
  const references: MysqlTableReference[] = [];

  for (const match of sql.matchAll(TABLE_REF_RE)) {
    const rawTableRef = match[1].replace(/\s+/g, '');
    const tableName = normalizeTableName(rawTableRef);
    if (tableName) {
      tables.push(tableName);
      const afterTable = sql.slice((match.index ?? 0) + match[0].length);
      const aliasMatch = /^\s+(?:as\s+)?(`[^`]+`|[A-Za-z_][\w$]*)/i.exec(afterTable);
      const alias = aliasMatch ? normalizeIdentifier(aliasMatch[1]) : undefined;
      references.push({
        table: tableName,
        schemaQualified: rawTableRef.includes('.'),
        ...(alias && !RESERVED_ALIAS_WORDS.has(alias) ? { alias } : {}),
      });
    }
  }

  return references;
}

function hasImplicitJoin(sql: string) {
  const fromMatch = /\bfrom\b/i.exec(sql);
  if (!fromMatch) {
    return false;
  }

  const afterFrom = sql.slice(fromMatch.index + fromMatch[0].length);
  const boundary = CLAUSE_BOUNDARY_RE.exec(afterFrom);
  const tableSegment = boundary ? afterFrom.slice(0, boundary.index) : afterFrom;

  return tableSegment.includes(',');
}

function withoutStringLiterals(sql: string) {
  let result = '';
  let quote: string | null = null;

  for (let index = 0; index < sql.length; index += 1) {
    const char = sql[index];
    const nextChar = sql[index + 1];

    if (quote) {
      if (char === '\\' && nextChar) {
        result += '  ';
        index += 1;
        continue;
      }
      if (char === quote) {
        if (nextChar === quote) {
          result += '  ';
          index += 1;
          continue;
        }
        quote = null;
      }
      result += ' ';
      continue;
    }

    if (char === '\'' || char === '"') {
      quote = char;
      result += ' ';
      continue;
    }

    result += char;
  }

  return result;
}

function isIdentifierChar(char: string | undefined) {
  return !!char && /[A-Za-z0-9_$]/.test(char);
}

function topLevelKeywordIndex(sql: string, keyword: string): number {
  const lowerSql = withoutStringLiterals(sql).toLowerCase();
  const lowerKeyword = keyword.toLowerCase();
  let depth = 0;
  let inBacktick = false;

  for (let index = 0; index < lowerSql.length; index += 1) {
    const char = lowerSql[index];
    if (char === '`') {
      inBacktick = !inBacktick;
      continue;
    }
    if (inBacktick) {
      continue;
    }
    if (char === '(') {
      depth += 1;
      continue;
    }
    if (char === ')') {
      depth = Math.max(0, depth - 1);
      continue;
    }
    if (
      depth === 0 &&
      lowerSql.startsWith(lowerKeyword, index) &&
      !isIdentifierChar(lowerSql[index - 1]) &&
      !isIdentifierChar(lowerSql[index + lowerKeyword.length])
    ) {
      return index;
    }
  }

  return -1;
}

function hasTopLevelKeyword(sql: string, keyword: string) {
  return topLevelKeywordIndex(sql, keyword) >= 0;
}

function joinSegmentsHaveConditions(sql: string) {
  const joinRe = /\b(?:inner|left|right|full|cross)?\s*join\b/gi;
  const joins = [...sql.matchAll(joinRe)];

  for (let index = 0; index < joins.length; index += 1) {
    const start = (joins[index].index ?? 0) + joins[index][0].length;
    const nextJoinIndex = index + 1 < joins.length ? joins[index + 1].index ?? sql.length : sql.length;
    const boundary = CLAUSE_BOUNDARY_RE.exec(sql.slice(start));
    const end = Math.min(nextJoinIndex, boundary ? start + boundary.index : sql.length);
    const segment = sql.slice(start, end);

    if (!/\b(on|using)\b/i.test(segment)) {
      return false;
    }
  }

  return true;
}

function getExplicitLimit(sql: string) {
  const limitIndex = topLevelKeywordIndex(sql, 'limit');
  if (limitIndex < 0) {
    return null;
  }

  const limitSql = sql.slice(limitIndex);
  const commaLimit = /\blimit\s+\d+\s*,\s*(\d+)\b/i.exec(limitSql);
  if (commaLimit) {
    return Number(commaLimit[1]);
  }

  const limit = /\blimit\s+(\d+)\b/i.exec(limitSql);
  if (limit) {
    return Number(limit[1]);
  }

  return null;
}

function hasOffsetLimit(sql: string) {
  const limitIndex = topLevelKeywordIndex(sql, 'limit');
  if (limitIndex < 0) {
    return false;
  }
  const limitSql = sql.slice(limitIndex);
  return /\blimit\s+\d+\s*,\s*\d+\b/i.test(limitSql) || /\blimit\s+\d+\s+offset\s+\d+\b/i.test(limitSql);
}

function tablesWithMaskedColumns(maskedColumns: MysqlSqlGuardOptions['maskedColumns']) {
  return new Set(
    Object.keys(maskedColumns ?? {})
      .map(key => key.toLowerCase().split('.')[0])
      .filter(Boolean)
  );
}

function hasMaskedWildcardProjection(
  sql: string,
  references: MysqlTableReference[],
  maskedColumns: MysqlSqlGuardOptions['maskedColumns']
) {
  const maskedTables = tablesWithMaskedColumns(maskedColumns);
  const referencedMaskedTables = references.filter(ref => maskedTables.has(ref.table));
  if (referencedMaskedTables.length === 0) {
    return false;
  }

  if (/^\s*select\s+\*/i.test(sql)) {
    return true;
  }

  for (const ref of referencedMaskedTables) {
    const wildcardOwners = [ref.table, ref.alias].filter(Boolean) as string[];
    for (const owner of wildcardOwners) {
      const wildcardRe = new RegExp(`(?:\\b${escapeRegExp(owner)}\\b|\`${escapeRegExp(owner)}\`)\\s*\\.\\s*\\*`, 'i');
      if (wildcardRe.test(sql)) {
        return true;
      }
    }
  }

  return false;
}

function hasMaskedColumnReference(
  sql: string,
  references: MysqlTableReference[],
  maskedColumns: MysqlSqlGuardOptions['maskedColumns']
) {
  if (!maskedColumns) {
    return false;
  }

  const literalSafeSql = withoutStringLiterals(sql);
  const referencedTables = new Set(references.map(ref => ref.table));
  for (const key of Object.keys(maskedColumns)) {
    const [rawTable, rawColumn] = key.toLowerCase().split('.');
    if (!rawTable || !rawColumn || !referencedTables.has(rawTable)) {
      continue;
    }

    const maskedOwners = references
      .filter(ref => ref.table === rawTable)
      .flatMap(ref => [ref.table, ref.alias].filter(Boolean) as string[]);
    if (maskedOwners.some(owner => qualifiedColumnRe(owner, rawColumn).test(literalSafeSql))) {
      return true;
    }

    const withoutQualifiedRefs = literalSafeSql.replace(
      /(?:`[^`]+`|[A-Za-z_][\w$]*)\s*\.\s*(?:`[^`]+`|[A-Za-z_][\w$]*)/g,
      ' '
    );
    const unqualifiedColumnRe = new RegExp(`(?:\\b${escapeRegExp(rawColumn)}\\b|\`${escapeRegExp(rawColumn)}\`)`, 'i');
    if (unqualifiedColumnRe.test(withoutQualifiedRefs)) {
      return true;
    }
  }

  return false;
}

function hasDerivedTable(sql: string) {
  return /\b(from|join)\s*\(/i.test(withoutStringLiterals(sql));
}

function escapeRegExp(value: string) {
  return value.replace(/[.*+?^${}()|[\]\\]/g, '\\$&');
}

function qualifiedColumnRe(owner: string, column: string) {
  return new RegExp(
    `(?:\\b${escapeRegExp(owner)}\\b|\`${escapeRegExp(owner)}\`)\\s*\\.\\s*(?:\\b${escapeRegExp(column)}\\b|\`${escapeRegExp(column)}\`)`,
    'i'
  );
}

export function validateMysqlReadOnlySql(
  sql: string,
  options: MysqlSqlGuardOptions = {}
): MysqlSqlGuardResult {
  let trimmedSql = sql.trim();

  if (!trimmedSql) {
    return deny('sql_empty');
  }

  if (/;\s*$/.test(trimmedSql)) {
    trimmedSql = trimmedSql.replace(/;\s*$/, '').trim();
  }

  if (COMMENT_RE.test(trimmedSql)) {
    return deny('comments_denied');
  }

  if (trimmedSql.includes(';')) {
    return deny('multiple_statements_denied');
  }

  if (METADATA_SQL_RE.test(trimmedSql)) {
    return deny('metadata_sql_denied');
  }

  if (/^with\b/i.test(trimmedSql)) {
    return deny('cte_denied');
  }

  if (DML_SQL_RE.test(trimmedSql)) {
    return deny('dml_sql_denied');
  }

  if (DDL_SQL_RE.test(trimmedSql)) {
    return deny('ddl_sql_denied');
  }

  if (DCL_SQL_RE.test(trimmedSql)) {
    return deny('dcl_sql_denied');
  }

  if (TRANSACTION_SQL_RE.test(trimmedSql) || /^call\b/i.test(trimmedSql)) {
    return deny('transaction_sql_denied');
  }

  if (!/^select\b/i.test(trimmedSql)) {
    return deny('only_select_allowed');
  }

  if (/\bunion\b/i.test(trimmedSql)) {
    return deny('union_denied');
  }

  if (/@[A-Za-z_][\w$]*/.test(trimmedSql)) {
    return deny('user_variable_denied');
  }

  if (/\btemporary\b/i.test(trimmedSql)) {
    return deny('temporary_table_denied');
  }

  if (DANGEROUS_FUNCTION_RE.test(trimmedSql)) {
    return deny('dangerous_function_denied');
  }

  if (/\binto\s+(out|dump)file\b/i.test(trimmedSql)) {
    return deny('file_output_denied');
  }

  if (hasImplicitJoin(trimmedSql)) {
    return deny('implicit_join_denied');
  }

  const tableReferences = extractMysqlTableReferences(trimmedSql);
  const tables = unique(tableReferences.map(ref => ref.table));
  const allowedTables = new Set((options.allowedTables ?? []).map(table => table.toLowerCase()));
  const blockedTables = new Set((options.blockedTables ?? []).map(table => table.toLowerCase()));

  if (tableReferences.some(ref => ref.schemaQualified)) {
    return deny('schema_qualified_table_denied', tables);
  }

  if (allowedTables.size === 0 || tables.length === 0 || tables.some(table => !allowedTables.has(table))) {
    return deny('table_not_allowed', tables);
  }

  if (tables.some(table => blockedTables.has(table))) {
    return deny('table_blocked', tables);
  }

  const maxJoinTables = clampNumber(options.maxJoinTables, 4, 1, 6);
  if (tables.length > maxJoinTables) {
    return deny('too_many_join_tables', tables);
  }

  if (!joinSegmentsHaveConditions(trimmedSql)) {
    return deny('join_condition_required', tables);
  }

  if (hasDerivedTable(trimmedSql)) {
    return deny('derived_table_denied', tables);
  }

  if (hasOffsetLimit(trimmedSql)) {
    return deny('offset_denied', tables);
  }

  if (LOCKING_READ_RE.test(withoutStringLiterals(trimmedSql))) {
    return deny('locking_read_denied', tables);
  }

  const maxRows = clampNumber(options.maxRows, 200, 1, 1000);
  const explicitLimit = getExplicitLimit(trimmedSql);
  if (explicitLimit !== null && explicitLimit > maxRows) {
    return deny('limit_exceeds_max_rows', tables);
  }

  if (
    hasMaskedWildcardProjection(trimmedSql, tableReferences, options.maskedColumns) ||
    hasMaskedColumnReference(trimmedSql, tableReferences, options.maskedColumns)
  ) {
    return deny('masked_column_denied', tables);
  }

  const appendedLimit = explicitLimit === null;
  return {
    ok: true,
    tables,
    limitedSql: appendedLimit
      ? `SELECT * FROM (${trimmedSql}) AS neta_limited_query LIMIT ?`
      : trimmedSql,
    appendedLimit,
  };
}

@Provide()
@Scope(ScopeEnum.Singleton)
export class MysqlQueryService {
  @Inject()
  mysqlPoolManager: MysqlPoolManager;

  @InjectEntityModel(NetaClawDataSourceQueryAuditEntity)
  auditRepo: Repository<NetaClawDataSourceQueryAuditEntity>;

  async executeReadOnly(params: ExecuteReadOnlyParams): Promise<MysqlReadOnlyQueryResult> {
    const startedAt = process.hrtime.bigint();
    const extra = params.source.extra ?? {};
    const maxRows = clampNumber(extra.maxRows, 200, 1, 1000);
    const guard = validateMysqlReadOnlySql(params.sql, {
      allowedTables: extra.allowedTables,
      blockedTables: extra.blockedTables,
      maxJoinTables: extra.maxJoinTables,
      maxRows,
      maskedColumns: extra.maskedColumns,
    });

    if (!guard.ok) {
      await this.writeAuditSafe(params, {
        status: 'rejected',
        rejectReason: guard.reason ?? 'unknown',
        elapsedMs: elapsedSince(startedAt),
        rowCount: 0,
        errorCode: null,
      });
      throw new Error(`mysql_sql_rejected: ${guard.reason ?? 'unknown'}`);
    }

    try {
      const pool = await this.mysqlPoolManager.getPool(params.source);
      const querySql = guard.limitedSql ?? params.sql;
      const [rawRows, fields] = await pool.execute(
        {
          sql: querySql,
          timeout: clampNumber(extra.queryTimeoutMs, 10000, 1000, 30000),
        },
        guard.appendedLimit ? [maxRows + 1] : []
      );
      const fetchedRows = Array.isArray(rawRows) ? rawRows as Record<string, unknown>[] : [];
      const truncated = !!guard.appendedLimit && fetchedRows.length > maxRows;
      const rows = truncated ? fetchedRows.slice(0, maxRows) : fetchedRows;
      const result: MysqlReadOnlyQueryResult = {
        columns: this.resolveColumns(fields, rows),
        rows,
        rowCount: rows.length,
        truncated,
        elapsedMs: elapsedSince(startedAt),
        sql: params.sql,
      };
      await this.writeAuditSafe(params, {
        status: 'success',
        rejectReason: null,
        elapsedMs: result.elapsedMs,
        rowCount: result.rowCount,
        errorCode: null,
      });
      return result;
    } catch (error) {
      const errorCode = this.sanitizeMysqlErrorCode(error);
      await this.writeAuditSafe(params, {
        status: 'failed',
        rejectReason: null,
        elapsedMs: elapsedSince(startedAt),
        rowCount: 0,
        errorCode,
      });
      throw new Error(errorCode ? `mysql_query_failed:${errorCode}` : 'mysql_query_failed');
    }
  }

  private resolveColumns(fields: unknown, rows: Record<string, unknown>[]): string[] {
    if (Array.isArray(fields) && fields.length > 0) {
      return fields
        .map(field => typeof field === 'object' && field && 'name' in field ? String((field as { name: unknown }).name) : '')
        .filter(Boolean);
    }
    return rows[0] ? Object.keys(rows[0]) : [];
  }

  private async writeAudit(
    params: ExecuteReadOnlyParams,
    audit: Pick<NetaClawDataSourceQueryAuditEntity, 'status' | 'rejectReason' | 'elapsedMs' | 'rowCount' | 'errorCode'>
  ): Promise<void> {
      await this.auditRepo.save({
      dataSourceId: params.source.id,
      agentId: params.agentId ?? null,
      userId: params.userId ?? null,
      toolCallId: params.toolCallId ?? null,
      sqlHash: createHash('sha256').update(params.sql).digest('hex'),
      sqlPreview: params.sql.slice(0, 1000),
      status: audit.status,
      rejectReason: audit.rejectReason,
      elapsedMs: audit.elapsedMs,
      rowCount: audit.rowCount,
      errorCode: audit.errorCode,
    });
  }

  private async writeAuditSafe(
    params: ExecuteReadOnlyParams,
    audit: Pick<NetaClawDataSourceQueryAuditEntity, 'status' | 'rejectReason' | 'elapsedMs' | 'rowCount' | 'errorCode'>
  ): Promise<void> {
    try {
      await this.writeAudit(params, audit);
    } catch {
      // Audit is best-effort and must not affect query or rejection behavior.
    }
  }

  private sanitizeMysqlErrorCode(error: unknown): string | null {
    const rawCode = typeof error === 'object' && error && 'code' in error
      ? String((error as { code?: unknown }).code)
      : '';
    const code = rawCode.replace(/[^A-Za-z0-9_]+/g, '_').slice(0, 50);
    return code || null;
  }
}
