"""Tests for Annie SDK connectors."""

from __future__ import annotations

import pytest

from annie_sdk import (
    ConnectorRelationship,
    ConnectorTable,
    ConnectorTableColumn,
    MockConnector,
    ValidationError,
)


# =========================================================================
# Pydantic models
# =========================================================================


class TestConnectorTableColumn:
    def test_minimal(self):
        col = ConnectorTableColumn(name="id")
        assert col.name == "id"
        assert col.type is None
        assert col.description is None

    def test_full(self):
        col = ConnectorTableColumn(name="status", type="string", description="Order status")
        assert col.name == "status"
        assert col.type == "string"
        assert col.description == "Order status"


class TestConnectorTable:
    def test_minimal(self):
        t = ConnectorTable(name="orders")
        assert t.name == "orders"
        assert t.description is None
        assert t.columns is None
        assert t.data_model is None

    def test_with_columns(self):
        cols = [ConnectorTableColumn(name="id", type="integer")]
        t = ConnectorTable(name="orders", columns=cols)
        assert len(t.columns) == 1
        assert t.columns[0].name == "id"

    def test_with_description(self):
        t = ConnectorTable(name="users", description="App users")
        assert t.description == "App users"


class TestConnectorRelationship:
    def test_fields(self):
        rel = ConnectorRelationship(
            primary_table="orders",
            primary_column="product_id",
            foreign_table="products",
            foreign_column="id",
            description="FK to products",
        )
        assert rel.primary_table == "orders"
        assert rel.primary_column == "product_id"
        assert rel.foreign_table == "products"
        assert rel.foreign_column == "id"
        assert rel.description == "FK to products"

    def test_description_optional(self):
        rel = ConnectorRelationship(
            primary_table="a",
            primary_column="b",
            foreign_table="c",
            foreign_column="d",
        )
        assert rel.description is None


# =========================================================================
# MockConnector
# =========================================================================


class TestMockConnector:
    def test_dialect(self):
        c = MockConnector(
            dialect_name="mysql",
            tables=[ConnectorTable(name="t")],
        )
        assert c.dialect == "mysql"

    def test_default_dialect(self, single_table_connector):
        assert single_table_connector.dialect == "postgres"

    def test_connect_disconnect(self, single_table_connector):
        assert not single_table_connector.is_connected()
        single_table_connector.connect()
        assert single_table_connector.is_connected()
        single_table_connector.disconnect()
        assert not single_table_connector.is_connected()

    def test_execute_records_queries(self, mock_connector):
        mock_connector.execute("SELECT 1")
        mock_connector.execute("SELECT 2")
        assert len(mock_connector.executed_queries) == 2
        assert mock_connector.executed_queries[0] == "SELECT 1"
        assert mock_connector.executed_queries[1] == "SELECT 2"

    def test_last_query(self, mock_connector):
        assert mock_connector.last_query is None
        mock_connector.execute("SELECT 1")
        assert mock_connector.last_query == "SELECT 1"

    def test_set_results(self, mock_connector):
        new_results = [{"x": 1}]
        mock_connector.set_results(new_results)
        result = mock_connector.execute("SELECT x")
        assert result == new_results

    def test_discover_tables(self, mock_connector):
        tables = mock_connector.discover_tables()
        assert "orders" in tables
        assert "products" in tables

    def test_discover_columns(self, mock_connector):
        cols = mock_connector.discover_columns("orders")
        col_names = [c.name for c in cols]
        assert "id" in col_names
        assert "name" in col_names
        assert "amount" in col_names

    def test_discover_columns_unknown_table(self, mock_connector):
        cols = mock_connector.discover_columns("nonexistent")
        assert cols == []


# =========================================================================
# BaseConnector (tested via MockConnector)
# =========================================================================


class TestBaseConnector:
    def test_requires_tables(self):
        with pytest.raises(ValidationError):
            MockConnector(tables=[])

    def test_requires_tables_none(self):
        with pytest.raises(ValidationError):
            MockConnector(tables=None)

    def test_get_table_config_found(self, mock_connector):
        config = mock_connector.get_table_config("orders")
        assert config is not None
        assert config.name == "orders"

    def test_get_table_config_not_found(self, mock_connector):
        assert mock_connector.get_table_config("nonexistent") is None

    def test_tables_property(self, mock_connector):
        assert len(mock_connector.tables) == 2

    def test_relationships_property(self, mock_connector):
        assert len(mock_connector.relationships) == 1
        assert mock_connector.relationships[0].primary_table == "orders"

    def test_relationships_empty_default(self, single_table_connector):
        assert single_table_connector.relationships == []

    def test_get_schema_context_basic(self, single_table_connector):
        ctx = single_table_connector.get_schema_context()
        assert "tables" in ctx
        assert len(ctx["tables"]) == 1
        table = ctx["tables"][0]
        assert table["name"] == "orders"
        assert "columns" in table
        assert len(table["columns"]) > 0

    def test_get_schema_context_with_relationships(self, mock_connector):
        ctx = mock_connector.get_schema_context()
        assert "relationships" in ctx
        assert len(ctx["relationships"]) == 1
        rel = ctx["relationships"][0]
        assert rel["primary_table"] == "orders"
        assert rel["foreign_table"] == "products"

    def test_get_schema_context_no_relationships(self, single_table_connector):
        ctx = single_table_connector.get_schema_context()
        assert "relationships" not in ctx

    def test_get_schema_context_includes_description(self, mock_connector):
        ctx = mock_connector.get_schema_context()
        orders = next(t for t in ctx["tables"] if t["name"] == "orders")
        assert orders["description"] == "Customer orders"

    def test_get_schema_context_column_description(self, mock_connector):
        ctx = mock_connector.get_schema_context()
        orders = next(t for t in ctx["tables"] if t["name"] == "orders")
        status_col = next(c for c in orders["columns"] if c["name"] == "status")
        assert status_col["description"] == "Order status"

    def test_get_schema_context_excludes_none_fields(self, mock_connector):
        ctx = mock_connector.get_schema_context()
        orders = next(t for t in ctx["tables"] if t["name"] == "orders")
        id_col = next(c for c in orders["columns"] if c["name"] == "id")
        # description is None, should not be in the dict
        assert "description" not in id_col

    def test_get_table_schema(self, mock_connector):
        schema = mock_connector.get_table_schema("orders")
        col_names = [c.name for c in schema]
        assert "id" in col_names

    def test_context_manager(self):
        connector = MockConnector(
            tables=[ConnectorTable(name="t")],
        )
        assert not connector.is_connected()
        with connector:
            assert connector.is_connected()
        assert not connector.is_connected()
