import { NetaClawDataSourceEntity } from '../src/modules/netaclaw/entity/data_source.js';
import { MysqlQueryService } from '../src/modules/netaclaw/service/mysql_query.js';

function makeSource(overrides: Partial<NetaClawDataSourceEntity> = {}) {
  return Object.assign(new NetaClawDataSourceEntity(), {
    id: 42,
    name: 'orders',
    type: 'mysql',
    database: 'shop',
    extra: {
      allowedTables: ['customers', 'orders'],
      blockedTables: [],
      maxRows: 2,
      maxJoinTables: 4,
      maskedColumns: { 'customers.phone': 'partial' },
      queryTimeoutMs: 5000,
    },
    ...overrides,
  });
}

function makeService(pool: any, auditRepo: any) {
  const service = new MysqlQueryService();
  service.mysqlPoolManager = {
    getPool: jest.fn().mockResolvedValue(pool),
  } as any;
  service.auditRepo = auditRepo;
  return service;
}

describe('MysqlQueryService', () => {
  it('executes guarded query and writes success audit', async () => {
    const execute = jest.fn().mockResolvedValue([
      [{ id: 1, name: 'Ada' }],
      [{ name: 'id' }, { name: 'name' }],
    ]);
    const auditRepo = { save: jest.fn().mockResolvedValue({}) };
    const service = makeService({ execute }, auditRepo);

    const result = await service.executeReadOnly({
      source: makeSource(),
      sql: 'SELECT id, name FROM customers LIMIT 1',
      agentId: 7,
      userId: 8,
      toolCallId: 'tool-1',
    });

    expect(execute).toHaveBeenCalledWith(
      { sql: 'SELECT id, name FROM customers LIMIT 1', timeout: 5000 },
      []
    );
    expect(result).toMatchObject({
      columns: ['id', 'name'],
      rows: [{ id: 1, name: 'Ada' }],
      rowCount: 1,
      truncated: false,
      sql: 'SELECT id, name FROM customers LIMIT 1',
    });
    expect(result.elapsedMs).toEqual(expect.any(Number));
    expect(auditRepo.save).toHaveBeenCalledWith(expect.objectContaining({
      dataSourceId: 42,
      agentId: 7,
      userId: 8,
      toolCallId: 'tool-1',
      status: 'success',
      rejectReason: null,
      rowCount: 1,
      errorCode: null,
      sqlPreview: 'SELECT id, name FROM customers LIMIT 1',
    }));
    expect(auditRepo.save.mock.calls[0][0].sqlHash).toMatch(/^[a-f0-9]{64}$/);
  });

  it('rejects masked column before querying MySQL and audits rejected', async () => {
    const execute = jest.fn();
    const auditRepo = { save: jest.fn().mockResolvedValue({}) };
    const service = makeService({ execute }, auditRepo);

    await expect(service.executeReadOnly({
      source: makeSource(),
      sql: 'SELECT phone FROM customers LIMIT 1',
    })).rejects.toThrow('mysql_sql_rejected: masked_column_denied');

    expect(execute).not.toHaveBeenCalled();
    expect(auditRepo.save).toHaveBeenCalledWith(expect.objectContaining({
      status: 'rejected',
      rejectReason: 'masked_column_denied',
      rowCount: 0,
      errorCode: null,
    }));
  });

  it('audits rejected SQL', async () => {
    const execute = jest.fn();
    const auditRepo = { save: jest.fn().mockResolvedValue({}) };
    const service = makeService({ execute }, auditRepo);

    await expect(service.executeReadOnly({
      source: makeSource(),
      sql: 'DELETE FROM customers',
      agentId: null,
      userId: null,
      toolCallId: null,
    })).rejects.toThrow('mysql_sql_rejected: dml_sql_denied');

    expect(execute).not.toHaveBeenCalled();
    expect(auditRepo.save).toHaveBeenCalledWith(expect.objectContaining({
      dataSourceId: 42,
      agentId: null,
      userId: null,
      toolCallId: null,
      status: 'rejected',
      rejectReason: 'dml_sql_denied',
      rowCount: 0,
    }));
  });

  it('uses no-limit wrapper params, truncates node-side, and audits fetched row count', async () => {
    const execute = jest.fn().mockResolvedValue([
      [{ id: 1 }, { id: 2 }, { id: 3 }],
      [{ name: 'id' }],
    ]);
    const auditRepo = { save: jest.fn().mockResolvedValue({}) };
    const service = makeService({ execute }, auditRepo);

    const result = await service.executeReadOnly({
      source: makeSource(),
      sql: 'SELECT id FROM customers',
    });

    expect(execute).toHaveBeenCalledWith(
      {
        sql: 'SELECT * FROM (SELECT id FROM customers) AS neta_limited_query LIMIT ?',
        timeout: 5000,
      },
      [3]
    );
    expect(result.rows).toEqual([{ id: 1 }, { id: 2 }]);
    expect(result.rowCount).toBe(2);
    expect(result.truncated).toBe(true);
    expect(auditRepo.save).toHaveBeenCalledWith(expect.objectContaining({
      status: 'success',
      rowCount: 2,
    }));
  });

  it('audits mysql failures with sanitized error codes', async () => {
    const execute = jest.fn().mockRejectedValue({ code: 'ER_PARSE_ERROR: detail' });
    const auditRepo = { save: jest.fn().mockResolvedValue({}) };
    const service = makeService({ execute }, auditRepo);

    await expect(service.executeReadOnly({
      source: makeSource(),
      sql: 'SELECT id FROM customers LIMIT 1',
    })).rejects.toThrow('mysql_query_failed:ER_PARSE_ERROR_detail');

    expect(auditRepo.save).toHaveBeenCalledWith(expect.objectContaining({
      status: 'failed',
      rejectReason: null,
      rowCount: 0,
      errorCode: 'ER_PARSE_ERROR_detail',
    }));
  });

  it('returns successful query results when success audit write fails', async () => {
    const execute = jest.fn().mockResolvedValue([
      [{ id: 1 }],
      [{ name: 'id' }],
    ]);
    const auditRepo = { save: jest.fn().mockRejectedValue(new Error('audit_down')) };
    const service = makeService({ execute }, auditRepo);

    await expect(service.executeReadOnly({
      source: makeSource(),
      sql: 'SELECT id FROM customers LIMIT 1',
    })).resolves.toMatchObject({
      rows: [{ id: 1 }],
      rowCount: 1,
      truncated: false,
    });
  });

  it('preserves guard rejection reason when rejected audit write fails', async () => {
    const execute = jest.fn();
    const auditRepo = { save: jest.fn().mockRejectedValue(new Error('audit_down')) };
    const service = makeService({ execute }, auditRepo);

    await expect(service.executeReadOnly({
      source: makeSource(),
      sql: 'DELETE FROM customers',
    })).rejects.toThrow('mysql_sql_rejected: dml_sql_denied');

    expect(execute).not.toHaveBeenCalled();
  });

  it('does not pass wrapper limit params for explicit LIMIT SQL containing a question mark literal', async () => {
    const execute = jest.fn().mockResolvedValue([
      [{ id: 1 }],
      [{ name: 'id' }],
    ]);
    const auditRepo = { save: jest.fn().mockResolvedValue({}) };
    const service = makeService({ execute }, auditRepo);

    await service.executeReadOnly({
      source: makeSource(),
      sql: 'SELECT id FROM customers WHERE name = "?" LIMIT 1',
    });

    expect(execute).toHaveBeenCalledWith(
      { sql: 'SELECT id FROM customers WHERE name = "?" LIMIT 1', timeout: 5000 },
      []
    );
  });

  it('does not pass wrapper limit params when explicit LIMIT SQL has surrounding whitespace', async () => {
    const execute = jest.fn().mockResolvedValue([
      [{ id: 1 }],
      [{ name: 'id' }],
    ]);
    const auditRepo = { save: jest.fn().mockResolvedValue({}) };
    const service = makeService({ execute }, auditRepo);

    await service.executeReadOnly({
      source: makeSource(),
      sql: ' SELECT id FROM customers LIMIT 1 ',
    });

    expect(execute).toHaveBeenCalledWith(
      { sql: 'SELECT id FROM customers LIMIT 1', timeout: 5000 },
      []
    );
  });

  it('passes wrapper limit params when LIMIT appears only in a string literal', async () => {
    const execute = jest.fn().mockResolvedValue([
      [{ id: 1 }, { id: 2 }, { id: 3 }],
      [{ name: 'id' }],
    ]);
    const auditRepo = { save: jest.fn().mockResolvedValue({}) };
    const service = makeService({ execute }, auditRepo);

    await service.executeReadOnly({
      source: makeSource(),
      sql: "SELECT id FROM customers WHERE note = 'limit 1'",
    });

    expect(execute).toHaveBeenCalledWith(
      {
        sql: "SELECT * FROM (SELECT id FROM customers WHERE note = 'limit 1') AS neta_limited_query LIMIT ?",
        timeout: 5000,
      },
      [3]
    );
  });
});