package tech.easyflow.datacenter.connector.impl;

import com.alibaba.fastjson2.JSONObject;
import org.junit.Test;
import tech.easyflow.datacenter.entity.DatacenterTable;
import tech.easyflow.datacenter.entity.DatacenterTableField;
import tech.easyflow.datacenter.meta.entity.DatacenterSource;
import tech.easyflow.common.web.exceptions.BusinessException;

import javax.sql.DataSource;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.ResultSetMetaData;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.List;

import static org.junit.Assert.assertTrue;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.atLeastOnce;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

/**
 * 项目 MySQL 连接器批量写入测试。
 */
public class ProjectMysqlConnectorBatchTest {

    /**
     * 验证原始 SQL 在单连接、单 ResultSet 中原样流式消费。
     *
     * @throws Exception JDBC 模拟初始化失败时抛出
     */
    @Test
    public void shouldStreamOriginalSqlWithoutPaginationRewrite()
            throws Exception {
        DataSource dataSource = mock(DataSource.class);
        Connection connection = mock(Connection.class);
        PreparedStatement statement =
                mock(PreparedStatement.class);
        ResultSet resultSet = mock(ResultSet.class);
        ResultSetMetaData metadata =
                mock(ResultSetMetaData.class);
        String sql =
                "SELECT id FROM sample LIMIT 10 FOR UPDATE";
        when(dataSource.getConnection())
                .thenReturn(connection);
        when(connection.prepareStatement(
                        eq(sql),
                        eq(ResultSet.TYPE_FORWARD_ONLY),
                        eq(ResultSet.CONCUR_READ_ONLY)))
                .thenReturn(statement);
        when(statement.executeQuery())
                .thenReturn(resultSet);
        when(resultSet.getMetaData())
                .thenReturn(metadata);
        when(metadata.getColumnCount()).thenReturn(1);
        when(metadata.getColumnLabel(1)).thenReturn("id");
        when(resultSet.next())
                .thenReturn(true, false);
        when(resultSet.getObject(1))
                .thenReturn(1L);
        List<String> ids = new ArrayList<>();

        new ProjectMysqlConnector(dataSource)
                .consumeBySql(
                        source(),
                        sql,
                        1_000,
                        row -> ids.add(
                                row.getString("id")));

        org.junit.Assert.assertEquals(
                List.of("1"), ids);
        verify(dataSource).getConnection();
        verify(connection).prepareStatement(
                sql,
                ResultSet.TYPE_FORWARD_ONLY,
                ResultSet.CONCUR_READ_ONLY);
        verify(statement).setFetchSize(
                Integer.MIN_VALUE);
        verify(statement).executeQuery();
    }

    /**
     * 验证多行写入仅获取一次连接，并按批次执行 JDBC batch。
     *
     * @throws Exception JDBC 模拟初始化失败时抛出
     */
    @Test
    public void shouldReuseSingleConnectionAndExecuteConfiguredBatches() throws Exception {
        DataSource dataSource = mock(DataSource.class);
        Connection connection = mock(Connection.class);
        List<PreparedStatement> statements = new ArrayList<>();
        when(dataSource.getConnection()).thenReturn(connection);
        when(connection.prepareStatement(anyString())).thenAnswer(invocation -> {
            PreparedStatement statement = mock(PreparedStatement.class);
            statements.add(statement);
            return statement;
        });

        ProjectMysqlConnector connector = new ProjectMysqlConnector(dataSource);
        DatacenterSource source = new DatacenterSource();
        source.setDatabaseName("easyflow");
        DatacenterTable table = new DatacenterTable();
        table.setTableName("sample");
        DatacenterTableField nameField = new DatacenterTableField();
        nameField.setFieldName("name");
        nameField.setWritable(1);
        table.setFields(List.of(nameField));

        List<JSONObject> rows = new ArrayList<>();
        for (int index = 0; index < 5; index++) {
            JSONObject row = new JSONObject();
            row.put("name", "row-" + index);
            rows.add(row);
        }

        connector.saveRows(source, table, rows, null, 2);

        verify(dataSource, times(1)).getConnection();
        if (statements.size() != 3) {
            throw new AssertionError("expected 3 JDBC batches but got " + statements.size());
        }
        int addBatchCalls = 0;
        for (PreparedStatement statement : statements) {
            verify(statement, times(1)).executeBatch();
            addBatchCalls += org.mockito.Mockito.mockingDetails(statement)
                    .getInvocations()
                    .stream()
                    .filter(invocation -> "addBatch".equals(invocation.getMethod().getName()))
                    .count();
        }
        if (addBatchCalls != rows.size()) {
            throw new AssertionError("expected " + rows.size() + " addBatch calls but got " + addBatchCalls);
        }
    }

    /**
     * 验证回执和业务批量写入在同一 JDBC 事务中提交。
     *
     * @throws Exception JDBC 模拟初始化失败时抛出
     */
    @Test
    public void shouldCommitReceiptAndRowsInSingleTransaction() throws Exception {
        DataSource dataSource = mock(DataSource.class);
        Connection connection = mock(Connection.class);
        PreparedStatement receiptStatement = mock(PreparedStatement.class);
        PreparedStatement queryStatement = mock(PreparedStatement.class);
        PreparedStatement rowStatement = mock(PreparedStatement.class);
        ResultSet resultSet = mock(ResultSet.class);
        when(dataSource.getConnection()).thenReturn(connection);
        when(connection.getAutoCommit()).thenReturn(true);
        when(connection.prepareStatement(anyString())).thenAnswer(invocation -> {
            String sql = invocation.getArgument(0);
            if (sql.startsWith("SELECT")) {
                return queryStatement;
            }
            return sql.contains("tb_datacenter_write_receipt")
                    ? receiptStatement
                    : rowStatement;
        });
        when(queryStatement.executeQuery()).thenReturn(resultSet);
        when(resultSet.next()).thenReturn(false);

        ProjectMysqlConnector connector = new ProjectMysqlConnector(dataSource);
        DatacenterSource source = source();
        DatacenterTable table = table();
        JSONObject row = new JSONObject();
        row.put("name", "row-1");

        boolean written = connector.saveRowsIdempotently(
                source,
                table,
                List.of(row),
                null,
                100,
                "receipt-key",
                "payload-hash");

        assertTrue(written);
        verify(connection).setAutoCommit(false);
        verify(receiptStatement).executeBatch();
        verify(receiptStatement).executeUpdate();
        verify(rowStatement).executeBatch();
        verify(connection, times(2)).commit();
        verify(connection, never()).rollback();
        verify(connection).setAutoCommit(true);
    }

    /**
     * 验证业务批量失败时回执与业务数据一并回滚。
     *
     * @throws Exception JDBC 模拟初始化失败时抛出
     */
    @Test
    public void shouldRollbackReceiptWhenBatchWriteFails() throws Exception {
        DataSource dataSource = mock(DataSource.class);
        Connection connection = mock(Connection.class);
        PreparedStatement receiptStatement = mock(PreparedStatement.class);
        PreparedStatement queryStatement = mock(PreparedStatement.class);
        PreparedStatement rowStatement = mock(PreparedStatement.class);
        ResultSet resultSet = mock(ResultSet.class);
        when(dataSource.getConnection()).thenReturn(connection);
        when(connection.getAutoCommit()).thenReturn(true);
        when(connection.prepareStatement(anyString())).thenAnswer(invocation -> {
            String sql = invocation.getArgument(0);
            if (sql.startsWith("SELECT")) {
                return queryStatement;
            }
            return sql.contains("tb_datacenter_write_receipt")
                    ? receiptStatement
                    : rowStatement;
        });
        when(queryStatement.executeQuery()).thenReturn(resultSet);
        when(resultSet.next()).thenReturn(false);
        when(rowStatement.executeBatch()).thenThrow(new SQLException("write failed"));

        ProjectMysqlConnector connector = new ProjectMysqlConnector(dataSource);
        JSONObject row = new JSONObject();
        row.put("name", "row-1");

        try {
            connector.saveRowsIdempotently(
                    source(),
                    table(),
                    List.of(row),
                    null,
                    100,
                    "receipt-key",
                    "payload-hash");
            throw new AssertionError("failed business batch must rollback");
        } catch (BusinessException expected) {
            assertTrue(expected.getMessage().contains("write failed"));
        }

        verify(connection, atLeastOnce()).rollback();
        verify(connection, never()).commit();
        verify(connection).setAutoCommit(true);
    }

    /**
     * 验证中间行失败时前序行已经提交，后续行不会执行。
     *
     * @throws Exception JDBC 模拟初始化失败时抛出
     */
    @Test
    public void shouldKeepEarlierRowsCommittedWhenMiddleRowFails() throws Exception {
        DataSource dataSource = mock(DataSource.class);
        Connection connection = mock(Connection.class);
        PreparedStatement receiptStatement = mock(PreparedStatement.class);
        PreparedStatement queryStatement = mock(PreparedStatement.class);
        PreparedStatement rowStatement = mock(PreparedStatement.class);
        ResultSet resultSet = mock(ResultSet.class);
        when(dataSource.getConnection()).thenReturn(connection);
        when(connection.getAutoCommit()).thenReturn(true);
        when(connection.prepareStatement(anyString())).thenAnswer(invocation -> {
            String sql = invocation.getArgument(0);
            if (sql.startsWith("SELECT")) {
                return queryStatement;
            }
            return sql.contains("tb_datacenter_write_receipt")
                    ? receiptStatement
                    : rowStatement;
        });
        when(queryStatement.executeQuery()).thenReturn(resultSet);
        when(resultSet.next()).thenReturn(false);
        when(rowStatement.executeBatch())
                .thenThrow(new SQLException("batch failed"))
                .thenReturn(new int[]{1})
                .thenThrow(new SQLException("middle row failed"));

        ProjectMysqlConnector connector = new ProjectMysqlConnector(dataSource);
        List<JSONObject> rows = List.of(
                row("row-0"), row("row-1"), row("row-2"));

        try {
            connector.saveRowsIdempotently(
                    source(),
                    table(),
                    rows,
                    null,
                    100,
                    "receipt-key",
                    "payload-hash");
            throw new AssertionError("middle row failure must be propagated");
        } catch (BusinessException expected) {
            assertTrue(expected.getMessage().contains("middle row failed"));
        }

        verify(dataSource, times(1)).getConnection();
        verify(connection, times(1)).commit();
        verify(connection, atLeastOnce()).rollback();
        verify(rowStatement, times(3)).executeBatch();
        verify(receiptStatement).executeBatch();
        verify(receiptStatement, times(2)).executeUpdate();
    }

    /**
     * 创建测试数据源元数据。
     *
     * @return 项目 MySQL 数据源
     */
    private DatacenterSource source() {
        DatacenterSource source = new DatacenterSource();
        source.setDatabaseName("easyflow");
        return source;
    }

    /**
     * 创建包含一个可写字段的测试表。
     *
     * @return 测试数据表
     */
    private DatacenterTable table() {
        DatacenterTable table = new DatacenterTable();
        table.setTableName("sample");
        DatacenterTableField nameField = new DatacenterTableField();
        nameField.setFieldName("name");
        nameField.setWritable(1);
        table.setFields(List.of(nameField));
        return table;
    }

    /**
     * 创建测试数据行。
     *
     * @param name 行名称
     * @return JSON 行
     */
    private JSONObject row(String name) {
        JSONObject row = new JSONObject();
        row.put("name", name);
        return row;
    }
}
