"""
Executor 模块测试
测试 SQLExecutor 和 BatchExecutor
"""
import pytest
from unittest.mock import Mock, MagicMock, patch
from sqlalchemy import MetaData

from xll_kit.database import (
    SQLExecutor,
    BatchExecutor,
    SQLResult,
)


class TestSQLExecutor:
    """SQL Executor 测试"""

    def test_init(self):
        """测试初始化"""
        executor = SQLExecutor(
            database_uri='mysql://user:pass@localhost/test',
            echo=True
        )

        assert executor.database_uri == 'mysql://user:pass@localhost/test'
        assert executor.echo == True

    def test_from_env(self, monkeypatch):
        """测试从环境变量创建"""
        monkeypatch.setenv('TEST_DB_URI', 'mysql://test@localhost/testdb')

        executor = SQLExecutor.from_env('TEST_DB_URI')

        assert 'mysql://test@localhost/testdb' in executor.database_uri

    def test_from_env_missing(self):
        """测试环境变量不存在时抛出异常"""
        with pytest.raises(ValueError, match="Environment variable .* not found"):
            SQLExecutor.from_env('NON_EXISTENT_VAR')

    def test_from_dict(self):
        """测试从字典创建"""
        config = {
            'database_uri': 'mysql://test@localhost/db',
            'echo': True
        }

        executor = SQLExecutor.from_dict(config)

        assert 'mysql://test@localhost/db' in executor.database_uri
        assert executor.echo == True

    def test_get_table_not_found(self):
        """测试获取不存在的表"""
        executor = SQLExecutor('sqlite:///:memory:')
        executor._metadata = MetaData()

        with pytest.raises(ValueError, match="Table .* not found"):
            executor.get_table('non_existent_table')

    @patch('xll_kit.database.executor.create_engine')
    def test_execute(self, mock_create_engine):
        """测试执行 SQL"""
        # 模拟引擎和连接
        mock_engine = MagicMock()
        mock_conn = MagicMock()
        mock_result = MagicMock()
        mock_result.rowcount = 5

        mock_conn.execute.return_value = mock_result
        # 正确配置上下文管理器
        mock_conn.__enter__ = Mock(return_value=mock_conn)
        mock_conn.__exit__ = Mock(return_value=False)
        mock_engine.connect.return_value = mock_conn
        mock_create_engine.return_value = mock_engine

        executor = SQLExecutor('mysql://test')
        sql_result = SQLResult(
            sql="INSERT INTO users (name) VALUES (:name)",
            params={'name': 'Test'}
        )

        affected = executor.execute(sql_result)

        assert affected == 5
        mock_conn.execute.assert_called_once()
        mock_conn.commit.assert_called_once()

    @patch('xll_kit.database.executor.create_engine')
    def test_execute_many(self, mock_create_engine):
        """测试批量执行多个 SQL"""
        mock_engine = MagicMock()
        mock_conn = MagicMock()
        mock_result = MagicMock()
        mock_result.rowcount = 2

        mock_conn.execute.return_value = mock_result
        # 正确配置上下文管理器
        mock_conn.__enter__ = Mock(return_value=mock_conn)
        mock_conn.__exit__ = Mock(return_value=False)
        mock_engine.connect.return_value = mock_conn
        mock_create_engine.return_value = mock_engine

        executor = SQLExecutor('mysql://test')
        sql_results = [
            SQLResult(sql="INSERT 1", params={}),
            SQLResult(sql="INSERT 2", params={}),
            SQLResult(sql="INSERT 3", params={}),
        ]

        total = executor.execute_many(sql_results)

        assert total == 6  # 3 * 2
        assert mock_conn.execute.call_count == 3

    @patch('xll_kit.database.executor.create_engine')
    def test_query(self, mock_create_engine):
        """测试查询"""
        mock_engine = MagicMock()
        mock_conn = MagicMock()
        mock_result = MagicMock()

        # 模拟查询结果
        mock_result.fetchall.return_value = [
            (1, 'Alice', 'alice@example.com', 25),
            (2, 'Bob', 'bob@example.com', 30)
        ]
        mock_result.keys.return_value = ['id', 'name', 'email', 'age']

        mock_conn.execute.return_value = mock_result
        # 正确配置上下文管理器
        mock_conn.__enter__ = Mock(return_value=mock_conn)
        mock_conn.__exit__ = Mock(return_value=False)
        mock_engine.connect.return_value = mock_conn
        mock_create_engine.return_value = mock_engine

        executor = SQLExecutor('mysql://test')
        results = executor.query(
            "SELECT * FROM users WHERE age > :min_age",
            {'min_age': 20}
        )

        assert len(results) == 2
        assert results[0]['name'] == 'Alice'
        assert results[0]['age'] == 25
        assert results[1]['name'] == 'Bob'

    @patch('xll_kit.database.executor.create_engine')
    def test_query_empty_result(self, mock_create_engine):
        """测试空查询结果"""
        mock_engine = MagicMock()
        mock_conn = MagicMock()
        mock_result = MagicMock()

        mock_result.fetchall.return_value = []
        mock_conn.execute.return_value = mock_result
        # 正确配置上下文管理器
        mock_conn.__enter__ = Mock(return_value=mock_conn)
        mock_conn.__exit__ = Mock(return_value=False)
        mock_engine.connect.return_value = mock_conn
        mock_create_engine.return_value = mock_engine

        executor = SQLExecutor('mysql://test')
        results = executor.query("SELECT * FROM users WHERE id = 999")

        assert results == []

    @patch('xll_kit.database.executor.create_engine')
    def test_query_one(self, mock_create_engine):
        """测试查询单条记录"""
        mock_engine = MagicMock()
        mock_conn = MagicMock()
        mock_result = MagicMock()

        mock_result.fetchall.return_value = [(1, 'Alice', 'alice@example.com')]
        mock_result.keys.return_value = ['id', 'name', 'email']

        mock_conn.execute.return_value = mock_result
        # 正确配置上下文管理器
        mock_conn.__enter__ = Mock(return_value=mock_conn)
        mock_conn.__exit__ = Mock(return_value=False)
        mock_engine.connect.return_value = mock_conn
        mock_create_engine.return_value = mock_engine

        executor = SQLExecutor('mysql://test')
        result = executor.query_one("SELECT * FROM users WHERE id = 1")

        assert result is not None
        assert result['name'] == 'Alice'
        assert result['email'] == 'alice@example.com'

    @patch('xll_kit.database.executor.create_engine')
    def test_query_one_no_result(self, mock_create_engine):
        """测试查询单条但无结果"""
        mock_engine = MagicMock()
        mock_conn = MagicMock()
        mock_result = MagicMock()

        mock_result.fetchall.return_value = []
        mock_conn.execute.return_value = mock_result
        # 正确配置上下文管理器
        mock_conn.__enter__ = Mock(return_value=mock_conn)
        mock_conn.__exit__ = Mock(return_value=False)
        mock_engine.connect.return_value = mock_conn
        mock_create_engine.return_value = mock_engine

        executor = SQLExecutor('mysql://test')
        result = executor.query_one("SELECT * FROM users WHERE id = 999")

        assert result is None

    @patch('xll_kit.database.executor.create_engine')
    def test_close(self, mock_create_engine):
        """测试关闭连接"""
        mock_engine = MagicMock()
        mock_create_engine.return_value = mock_engine

        executor = SQLExecutor('mysql://test')
        # 触发引擎创建
        _ = executor.engine
        executor.close()

        mock_engine.dispose.assert_called_once()


class TestBatchExecutor:
    """Batch Executor 测试"""

    def test_init(self):
        """测试初始化"""
        mock_executor = MagicMock(spec=SQLExecutor)
        batch_executor = BatchExecutor(mock_executor, batch_size=100)

        assert batch_executor.executor == mock_executor
        assert batch_executor.batch_size == 100

    def test_default_batch_size(self):
        """测试默认批次大小"""
        mock_executor = MagicMock(spec=SQLExecutor)
        batch_executor = BatchExecutor(mock_executor)

        assert batch_executor.batch_size == 1000

    def test_execute_batches_single_batch(self):
        """测试单批执行"""
        mock_executor = MagicMock(spec=SQLExecutor)
        mock_executor.execute_many.return_value = 50

        batch_executor = BatchExecutor(mock_executor, batch_size=100)
        sql_results = [SQLResult(sql=f"INSERT {i}", params={}) for i in range(50)]

        total = batch_executor.execute_batches(sql_results)

        assert total == 50
        assert mock_executor.execute_many.call_count == 1

    def test_execute_batches_multiple_batches(self):
        """测试多批执行"""
        mock_executor = MagicMock(spec=SQLExecutor)
        mock_executor.execute_many.return_value = 100

        batch_executor = BatchExecutor(mock_executor, batch_size=100)
        sql_results = [SQLResult(sql=f"INSERT {i}", params={}) for i in range(250)]

        total = batch_executor.execute_batches(sql_results)

        # 250条数据，batch_size=100，应该分3批
        assert mock_executor.execute_many.call_count == 3
        assert total == 300  # 3批 * 100

    def test_execute_batches_exact_batch_size(self):
        """测试正好整批的数据"""
        mock_executor = MagicMock(spec=SQLExecutor)
        mock_executor.execute_many.return_value = 100

        batch_executor = BatchExecutor(mock_executor, batch_size=100)
        sql_results = [SQLResult(sql=f"INSERT {i}", params={}) for i in range(200)]

        total = batch_executor.execute_batches(sql_results)

        # 200条数据，batch_size=100，应该分2批
        assert mock_executor.execute_many.call_count == 2
        assert total == 200


class TestExecutorContextManager:
    """测试 Executor 的上下文管理器"""

    @patch('xll_kit.database.executor.create_engine')
    def test_get_connection_success(self, mock_create_engine):
        """测试成功的连接上下文"""
        mock_engine = MagicMock()
        mock_conn = MagicMock()

        # 正确配置上下文管理器
        mock_conn.__enter__ = Mock(return_value=mock_conn)
        mock_conn.__exit__ = Mock(return_value=False)
        mock_engine.connect.return_value = mock_conn
        mock_create_engine.return_value = mock_engine

        executor = SQLExecutor('mysql://test')

        with executor.get_connection() as conn:
            assert conn == mock_conn

        mock_conn.commit.assert_called_once()
        mock_conn.close.assert_called_once()

    @patch('xll_kit.database.executor.create_engine')
    def test_get_connection_rollback_on_error(self, mock_create_engine):
        """测试异常时回滚"""
        mock_engine = MagicMock()
        mock_conn = MagicMock()

        # 正确配置上下文管理器
        mock_conn.__enter__ = Mock(return_value=mock_conn)
        mock_conn.__exit__ = Mock(return_value=False)
        mock_engine.connect.return_value = mock_conn
        mock_create_engine.return_value = mock_engine

        executor = SQLExecutor('mysql://test')

        with pytest.raises(ValueError):
            with executor.get_connection() as conn:
                raise ValueError("Test error")

        mock_conn.rollback.assert_called_once()
        mock_conn.close.assert_called_once()
        mock_conn.commit.assert_not_called()