from sqlalchemy import UniqueConstraint, inspect, text
from sqlalchemy.ext.asyncio import AsyncEngine

from session_hub_core.models import Base, MessageModel, SessionModel, SyncLogModel

SESSION_COLUMNS = {
    "id",
    "host",
    "tool",
    "session_id",
    "title",
    "working_directory",
    "model_name",
    "created_at",
    "last_activity_at",
    "message_count",
    "hidden",
    "metadata",
    "synced_at",
    "content_tsv",
    "embedding",
    "embedding_model",
    "embedded_at",
}

MESSAGE_COLUMNS = {
    "id",
    "session_pk",
    "host",
    "tool",
    "session_id",
    "node_id",
    "parent_node_id",
    "message_id",
    "role",
    "content",
    "created_at",
    "metadata",
    "synced_at",
    "content_tsv",
    "embedding",
    "embedding_model",
    "embedded_at",
}

SYNC_LOG_COLUMNS = {
    "id",
    "host",
    "tool",
    "started_at",
    "finished_at",
    "sessions_upserted",
    "messages_upserted",
    "status",
    "error",
}


def test_tables_defined():
    assert set(Base.metadata.tables) == {"sessions", "messages", "sync_log"}


def test_sessions_columns():
    cols = SessionModel.__table__.columns
    assert set(cols.keys()) == SESSION_COLUMNS
    assert cols["id"].primary_key
    assert cols["host"].nullable is False
    assert cols["tool"].nullable is False
    assert cols["session_id"].nullable is False
    assert cols["title"].nullable is False
    assert cols["working_directory"].nullable is False
    assert cols["model_name"].nullable is True
    assert cols["created_at"].nullable is True
    assert cols["last_activity_at"].nullable is True
    assert cols["message_count"].nullable is False
    assert cols["hidden"].nullable is False
    assert cols["metadata"].nullable is False
    assert cols["synced_at"].nullable is False
    assert str(cols["metadata"].type) == "JSONB"
    assert str(cols["content_tsv"].type) == "TSVECTOR"
    assert str(cols["embedding"].type) == "VECTOR(1536)"


def test_messages_columns():
    cols = MessageModel.__table__.columns
    assert set(cols.keys()) == MESSAGE_COLUMNS
    assert cols["id"].primary_key
    assert cols["session_pk"].nullable is False
    assert cols["node_id"].nullable is False
    assert cols["parent_node_id"].nullable is True
    assert cols["message_id"].nullable is False
    assert cols["role"].nullable is False
    assert cols["content"].nullable is False
    assert str(cols["metadata"].type) == "JSONB"
    assert str(cols["content_tsv"].type) == "TSVECTOR"
    assert str(cols["embedding"].type) == "VECTOR(1536)"


def test_sync_log_columns():
    cols = SyncLogModel.__table__.columns
    assert set(cols.keys()) == SYNC_LOG_COLUMNS
    assert cols["id"].primary_key
    assert cols["host"].nullable is False
    assert cols["tool"].nullable is True
    assert cols["started_at"].nullable is False
    assert cols["finished_at"].nullable is True
    assert cols["status"].nullable is False
    assert cols["error"].nullable is True


def test_sessions_unique_constraint():
    constraints = [c for c in SessionModel.__table__.constraints if isinstance(c, UniqueConstraint)]
    assert len(constraints) == 1
    assert sorted(constraints[0].columns.keys()) == ["host", "session_id", "tool"]


def test_messages_unique_constraint():
    constraints = [c for c in MessageModel.__table__.constraints if isinstance(c, UniqueConstraint)]
    assert len(constraints) == 1
    assert sorted(constraints[0].columns.keys()) == ["host", "node_id", "session_id", "tool"]


def test_messages_foreign_key_cascade():
    fks = list(MessageModel.__table__.foreign_keys)
    assert len(fks) == 1
    assert fks[0].column.table.name == "sessions"
    assert fks[0].column.name == "id"
    assert fks[0].ondelete == "CASCADE"


def test_sessions_indexes():
    indexes = {i.name: i for i in SessionModel.__table__.indexes}
    assert "ix_sessions_tool" in indexes
    assert "ix_sessions_host" in indexes
    assert "ix_sessions_last_activity_at" in indexes
    assert "ix_sessions_working_directory" in indexes
    assert "idx_sessions_content_tsv" in indexes
    assert "idx_sessions_embedding" in indexes
    assert indexes["idx_sessions_content_tsv"].dialect_options["postgresql"]["using"] == "gin"
    assert indexes["idx_sessions_embedding"].dialect_options["postgresql"]["using"] == "hnsw"
    assert indexes["idx_sessions_embedding"].dialect_options["postgresql"]["ops"] == {
        "embedding": "vector_cosine_ops"
    }


def test_messages_indexes():
    indexes = {i.name: i for i in MessageModel.__table__.indexes}
    assert "ix_messages_session_pk" in indexes
    assert "ix_messages_role" in indexes
    assert "ix_messages_created_at" in indexes
    assert "ix_messages_host_tool_session_id" in indexes
    assert "idx_messages_content_tsv" in indexes
    assert "idx_messages_content_trgm" in indexes
    assert "idx_messages_embedding" in indexes
    assert indexes["idx_messages_content_tsv"].dialect_options["postgresql"]["using"] == "gin"
    assert indexes["idx_messages_content_trgm"].dialect_options["postgresql"]["ops"] == {
        "content": "gin_trgm_ops"
    }
    assert indexes["idx_messages_embedding"].dialect_options["postgresql"]["using"] == "hnsw"
    assert indexes["idx_messages_embedding"].dialect_options["postgresql"]["ops"] == {
        "embedding": "vector_cosine_ops"
    }


def test_sync_log_indexes():
    indexes = {i.name: i for i in SyncLogModel.__table__.indexes}
    assert "ix_sync_log_host_started_at" in indexes


async def test_db_schema_matches(engine: AsyncEngine) -> None:
    async with engine.connect() as conn:

        def _inspect(sync_conn):
            inspector = inspect(sync_conn)
            return {
                "tables": set(inspector.get_table_names()),
                "sessions_columns": {c["name"]: c for c in inspector.get_columns("sessions")},
                "messages_columns": {c["name"]: c for c in inspector.get_columns("messages")},
                "sync_log_columns": {c["name"]: c for c in inspector.get_columns("sync_log")},
                "sessions_indexes": {i["name"]: i for i in inspector.get_indexes("sessions")},
                "messages_indexes": {i["name"]: i for i in inspector.get_indexes("messages")},
                "session_uq": {
                    tuple(sorted(u["column_names"]))
                    for u in inspector.get_unique_constraints("sessions")
                },
                "message_uq": {
                    tuple(sorted(u["column_names"]))
                    for u in inspector.get_unique_constraints("messages")
                },
            }

        result = await conn.run_sync(_inspect)
    assert result["tables"] >= {"sessions", "messages", "sync_log"}
    sessions_columns = result["sessions_columns"]
    assert set(sessions_columns) == SESSION_COLUMNS
    assert sessions_columns["host"]["nullable"] is False
    assert sessions_columns["model_name"]["nullable"] is True
    assert sessions_columns["embedding"]["type"].compile() == "VECTOR(1536)"
    messages_columns = result["messages_columns"]
    assert set(messages_columns) == MESSAGE_COLUMNS
    sync_log_columns = result["sync_log_columns"]
    assert set(sync_log_columns) == SYNC_LOG_COLUMNS
    sessions_indexes = result["sessions_indexes"]
    assert "idx_sessions_content_tsv" in sessions_indexes
    assert "idx_sessions_embedding" in sessions_indexes
    assert (
        sessions_indexes["idx_sessions_embedding"]["dialect_options"]["postgresql_using"] == "hnsw"
    )
    messages_indexes = result["messages_indexes"]
    assert "idx_messages_content_trgm" in messages_indexes
    assert (
        messages_indexes["idx_messages_content_trgm"]["dialect_options"]["postgresql_using"]
        == "gin"
    )
    session_uq = result["session_uq"]
    assert ("host", "session_id", "tool") in session_uq
    message_uq = result["message_uq"]
    assert ("host", "node_id", "session_id", "tool") in message_uq


async def test_db_triggers_created(engine: AsyncEngine) -> None:
    async with engine.connect() as conn:
        functions = (
            (
                await conn.execute(
                    text(
                        "SELECT proname FROM pg_proc "
                        "WHERE proname IN ('sessions_content_tsv_update', "
                        "'messages_session_tsv_update', 'messages_content_tsv_update')"
                    )
                )
            )
            .scalars()
            .all()
        )
        triggers = (
            (
                await conn.execute(
                    text(
                        "SELECT tgname FROM pg_trigger WHERE NOT tgisinternal AND tgname IN ("
                        "'sessions_content_tsv_trigger', 'messages_session_tsv_trigger', "
                        "'messages_content_tsv_trigger')"
                    )
                )
            )
            .scalars()
            .all()
        )
    assert set(functions) == {
        "sessions_content_tsv_update",
        "messages_session_tsv_update",
        "messages_content_tsv_update",
    }
    assert set(triggers) == {
        "sessions_content_tsv_trigger",
        "messages_session_tsv_trigger",
        "messages_content_tsv_trigger",
    }
