import pytest
import polars as pl
from bear_lake import Database, connect, table


class TestQueryOperations:
    """Test query functionality."""

    def test_query_all_data(self, db, sample_schema, sample_data):
        """Test querying all data from a table."""
        db.create(
            name="users",
            schema=sample_schema,
            partition_keys=["city"],
            primary_keys=["id"],
        )

        db.insert("users", sample_data)

        # Query all data
        result = db.query(pl.scan_parquet(f"{db.path}/users/**/*.parquet"))
        assert len(result) == len(sample_data)

    def test_query_with_filter(self, db, sample_schema, sample_data):
        """Test querying with a filter condition."""
        db.create(
            name="users",
            schema=sample_schema,
            partition_keys=["city"],
            primary_keys=["id"],
        )

        db.insert("users", sample_data)

        # Query with filter
        result = db.query(
            pl.scan_parquet(f"{db.path}/users/**/*.parquet").filter(pl.col("age") > 30)
        )

        assert len(result) == 2  # Charlie (35) and Eve (32)
        assert all(result["age"] > 30)

    def test_query_with_select(self, db, sample_schema, sample_data):
        """Test querying with column selection."""
        db.create(
            name="users",
            schema=sample_schema,
            partition_keys=["city"],
            primary_keys=["id"],
        )

        db.insert("users", sample_data)

        # Query with select
        result = db.query(
            pl.scan_parquet(f"{db.path}/users/**/*.parquet").select(["name", "city"])
        )

        assert result.columns == ["name", "city"]
        assert len(result) == len(sample_data)

    def test_query_with_aggregation(self, db, sample_schema, sample_data):
        """Test querying with aggregation."""
        db.create(
            name="users",
            schema=sample_schema,
            partition_keys=["city"],
            primary_keys=["id"],
        )

        db.insert("users", sample_data)

        # Query with aggregation - count by city
        result = db.query(
            pl.scan_parquet(f"{db.path}/users/**/*.parquet")
            .group_by("city")
            .agg(pl.col("id").count().alias("count"))
            .sort("city")
        )

        # NYC: 2, LA: 2, SF: 1
        assert len(result) == 3

    def test_query_with_join(self, db, sample_schema, sample_data):
        """Test querying with joins across tables."""
        # Create users table
        db.create(
            name="users",
            schema=sample_schema,
            partition_keys=["city"],
            primary_keys=["id"],
        )
        db.insert("users", sample_data)

        # Create orders table
        orders_schema = {
            "order_id": pl.Int64,
            "user_id": pl.Int64,
            "amount": pl.Float64,
        }

        orders_data = pl.DataFrame(
            {
                "order_id": [1, 2, 3],
                "user_id": [1, 2, 1],
                "amount": [100.0, 200.0, 150.0],
            }
        )

        db.create(
            name="orders",
            schema=orders_schema,
            partition_keys=["user_id"],  # Use user_id as partition key
            primary_keys=["order_id"],
        )
        db.insert("orders", orders_data)

        # Query with join
        users_lf = pl.scan_parquet(f"{db.path}/users/**/*.parquet")
        orders_lf = pl.scan_parquet(f"{db.path}/orders/**/*.parquet")

        result = db.query(
            users_lf.join(orders_lf, left_on="id", right_on="user_id").select(
                ["name", "order_id", "amount"]
            )
        )

        assert len(result) == 3  # 3 orders
        assert "name" in result.columns
        assert "order_id" in result.columns

    def test_query_empty_table(self, db, sample_schema):
        """Test querying an empty table."""
        db.create(
            name="users",
            schema=sample_schema,
            partition_keys=["city"],
            primary_keys=["id"],
        )

        # Query without any inserts - should handle gracefully
        # This will fail if no parquet files exist, which is expected
        # The table needs at least one partition to be queryable
        with pytest.raises(Exception):  # FileNotFoundError or similar
            db.query(pl.scan_parquet(f"{db.path}/users/**/*.parquet"))


class TestConnectFunction:
    """Test the connect() helper function."""

    def test_connect_returns_database(self, temp_db_path):
        """Test that connect returns a Database instance."""
        db = connect(temp_db_path)
        assert isinstance(db, Database)
        assert db.path == temp_db_path

    def test_connect_sets_global_state(self, temp_db_path):
        """Test that connect sets global state for table() function."""
        from bear_lake import DATABASE_PATH, CONNECTED

        db = connect(temp_db_path)

        # Import again to get updated global values
        import bear_lake

        assert bear_lake.DATABASE_PATH == temp_db_path
        assert bear_lake.CONNECTED is True


class TestTableFunction:
    """Test the table() helper function."""

    def test_table_function_basic(self, temp_db_path, sample_schema, sample_data):
        """Test the table() helper function."""
        # Connect first
        db = connect(temp_db_path)

        # Create and populate table
        db.create(
            name="users",
            schema=sample_schema,
            partition_keys=["city"],
            primary_keys=["id"],
        )
        db.insert("users", sample_data)

        # Use table() function
        lf = table("users")
        assert isinstance(lf, pl.LazyFrame)

        # Query using the lazy frame
        result = db.query(lf)
        assert len(result) == len(sample_data)

    def test_table_function_not_connected(self):
        """Test that table() raises error when not connected."""
        # Reset global state
        from bear_lake import connection

        # Save current state
        original_connected = connection.CONNECTED

        try:
            connection.CONNECTED = False

            with pytest.raises(RuntimeError, match="Not connected to database"):
                table("users")
        finally:
            # Restore original state
            connection.CONNECTED = original_connected

    def test_table_function_with_filter(self, temp_db_path, sample_schema, sample_data):
        """Test using table() function with filters."""
        db = connect(temp_db_path)

        db.create(
            name="users",
            schema=sample_schema,
            partition_keys=["city"],
            primary_keys=["id"],
        )
        db.insert("users", sample_data)

        # Use table() with filter
        result = db.query(table("users").filter(pl.col("city") == "NYC"))

        assert len(result) == 2
        assert all(result["city"] == "NYC")


class TestComplexQueries:
    """Test complex query scenarios."""

    def test_query_multiple_partitions(
        self, db, multi_partition_schema, partitioned_data
    ):
        """Test querying across multiple partition levels."""
        db.create(
            name="users",
            schema=multi_partition_schema,
            partition_keys=["country", "city"],
            primary_keys=["id"],
        )

        db.insert("users", partitioned_data)

        # Query specific partition
        result = db.query(
            pl.scan_parquet(f"{db.path}/users/**/*.parquet").filter(
                pl.col("city") == "NYC"
            )
        )

        nyc_count = len(partitioned_data.filter(pl.col("city") == "NYC"))
        assert len(result) == nyc_count

    def test_query_with_sorting(self, db, sample_schema, sample_data):
        """Test querying with sorting."""
        db.create(
            name="users",
            schema=sample_schema,
            partition_keys=["city"],
            primary_keys=["id"],
        )

        db.insert("users", sample_data)

        # Query with sorting
        result = db.query(
            pl.scan_parquet(f"{db.path}/users/**/*.parquet").sort(
                "age", descending=True
            )
        )

        ages = result["age"].to_list()
        assert ages == sorted(ages, reverse=True)

    def test_query_with_limit(self, db, sample_schema, sample_data):
        """Test querying with limit."""
        db.create(
            name="users",
            schema=sample_schema,
            partition_keys=["city"],
            primary_keys=["id"],
        )

        db.insert("users", sample_data)

        # Query with limit
        result = db.query(pl.scan_parquet(f"{db.path}/users/**/*.parquet").limit(3))

        assert len(result) == 3

    def test_query_with_expressions(self, db, sample_schema, sample_data):
        """Test querying with computed expressions."""
        db.create(
            name="users",
            schema=sample_schema,
            partition_keys=["city"],
            primary_keys=["id"],
        )

        db.insert("users", sample_data)

        # Query with expression - create age_group
        result = db.query(
            pl.scan_parquet(f"{db.path}/users/**/*.parquet").with_columns(
                (pl.col("age") / 10).cast(pl.Int32).alias("age_decade")
            )
        )

        assert "age_decade" in result.columns
        assert all(result["age_decade"].is_in([2, 3]))
