"""Tests for configurable table engine support."""

from aaiclick import create_object, create_object_from_value
from aaiclick.data.data_context import data_context, get_ch_client
from aaiclick.data.models import ColumnInfo, ENGINE_MEMORY, ENGINE_MERGE_TREE, Schema


async def test_context_with_memory_engine():
    """Test non-default engine (Memory) set on data_context."""
    async with data_context(engine=ENGINE_MEMORY):
        ch_client = get_ch_client()
        obj_a = await create_object_from_value([1, 2, 3])
        obj_b = await create_object_from_value([4, 5, 6])

        # Verify source tables use Memory engine
        result = await ch_client.query(f"""
            SELECT engine FROM system.tables WHERE name = '{obj_a.table}'
        """)
        assert result.result_rows[0][0] == "Memory"

        # Operator should also create Memory table
        obj_c = await (obj_a + obj_b)

        result = await ch_client.query(f"""
            SELECT engine FROM system.tables WHERE name = '{obj_c.table}'
        """)
        assert result.result_rows[0][0] == "Memory"

        # Verify data is correct
        assert await obj_c.data() == [5, 7, 9]


async def test_create_object_with_engine_override():
    """Test non-default engine passed directly to create_object."""
    async with data_context():  # Default MergeTree
        ch_client = get_ch_client()
        # Override with Memory engine
        schema = Schema(fieldtype="a", columns={"aai_id": ColumnInfo("UInt64"), "value": ColumnInfo("Int64")})
        obj_a = await create_object(schema, engine=ENGINE_MEMORY)

        result = await ch_client.query(f"""
            SELECT engine FROM system.tables WHERE name = '{obj_a.table}'
        """)
        assert result.result_rows[0][0] == "Memory"

        # Without override, should use context default (MergeTree)
        obj_b = await create_object(schema)

        result = await ch_client.query(f"""
            SELECT engine FROM system.tables WHERE name = '{obj_b.table}'
        """)
        assert result.result_rows[0][0] == "MergeTree"


async def test_mixed_engine_scenario():
    """Test mix of context engine and per-object override."""
    async with data_context(engine=ENGINE_MEMORY):
        ch_client = get_ch_client()
        # Create with context default (Memory)
        obj_a = await create_object_from_value([10, 20, 30])

        result = await ch_client.query(f"""
            SELECT engine FROM system.tables WHERE name = '{obj_a.table}'
        """)
        assert result.result_rows[0][0] == "Memory"

        # Override to MergeTree for specific object
        schema = Schema(fieldtype="a", columns={"aai_id": ColumnInfo("UInt64"), "value": ColumnInfo("Int64")})
        obj_b = await create_object(schema, engine=ENGINE_MERGE_TREE)

        result = await ch_client.query(f"""
            SELECT engine FROM system.tables WHERE name = '{obj_b.table}'
        """)
        assert result.result_rows[0][0] == "MergeTree"

        # Insert data into obj_b (aai_id generated by DEFAULT)
        await ch_client.insert(obj_b.table, [[100], [200], [300]], column_names=["value"])

        # Operator between Memory and MergeTree objects
        # Result should use context default (Memory)
        obj_c = await (obj_a + obj_b)

        result = await ch_client.query(f"""
            SELECT engine FROM system.tables WHERE name = '{obj_c.table}'
        """)
        assert result.result_rows[0][0] == "Memory"

        # Verify data is correct
        assert await obj_c.data() == [110, 220, 330]
