"""
Tests for migration 0003: Add research indexes.

Tests cover:
- Index creation on research_tasks and research_history
- Full migration chain from empty database
- Idempotency (running migrations multiple times)
- Downgrade behavior
- Data preservation during migration
- Edge cases (partial tables, empty tables, in-memory databases)
"""

from alembic import command
import pytest
from sqlalchemy import create_engine, inspect, text

from local_deep_research.database.alembic_runner import (
    get_alembic_config,
    get_current_revision,
    get_head_revision,
    needs_migration,
    run_migrations,
)

# Expected indexes from migration 0003
RESEARCH_TASKS_INDEXES = {
    "ix_research_tasks_status": ["status"],
    "ix_research_tasks_created_at": ["created_at"],
    "idx_research_task_status_created": ["status", "created_at"],
    "idx_research_task_priority_status": ["priority", "status"],
}

RESEARCH_HISTORY_INDEXES = {
    "ix_research_history_mode": ["mode"],
    "ix_research_history_status": ["status"],
    "ix_research_history_created_at": ["created_at"],
    "idx_research_history_status_created": ["status", "created_at"],
    "idx_research_history_mode_status": ["mode", "status"],
}


def _get_indexes_by_name(engine, table_name):
    """Get a dict of {index_name: [column_names]} for a table."""
    insp = inspect(engine)
    if not insp.has_table(table_name):
        return {}
    return {
        idx["name"]: idx["column_names"]
        for idx in insp.get_indexes(table_name)
        if idx["name"] is not None
    }


def _run_upgrade_to(engine, revision):
    """Run migrations up to a specific revision."""
    config = get_alembic_config(engine)
    with engine.begin() as conn:
        config.attributes["connection"] = conn
        command.upgrade(config, revision)


def _run_downgrade_to(engine, revision):
    """Run downgrade to a specific revision."""
    config = get_alembic_config(engine)
    with engine.begin() as conn:
        config.attributes["connection"] = conn
        command.downgrade(config, revision)


@pytest.fixture
def fresh_engine(tmp_path):
    """Create a fresh SQLite engine (empty database, no tables)."""
    db_path = tmp_path / "fresh_0003_test.db"
    engine = create_engine(f"sqlite:///{db_path}")
    yield engine
    engine.dispose()


@pytest.fixture
def migrated_to_0002_engine(tmp_path):
    """Create a database migrated to revision 0002 (tables exist, no research indexes)."""
    db_path = tmp_path / "migrated_0002_test.db"
    engine = create_engine(f"sqlite:///{db_path}")
    _run_upgrade_to(engine, "0002")
    yield engine
    engine.dispose()


@pytest.fixture
def fully_migrated_engine(tmp_path):
    """Create a database migrated to head."""
    db_path = tmp_path / "fully_migrated_0003_test.db"
    engine = create_engine(f"sqlite:///{db_path}")
    run_migrations(engine)
    yield engine
    engine.dispose()


class TestMigration0003UpgradeIndexes:
    """Tests that verify index creation on upgrade to 0003."""

    def test_creates_research_task_single_column_indexes(
        self, fully_migrated_engine
    ):
        """Verify ix_research_tasks_status and ix_research_tasks_created_at exist."""
        indexes = _get_indexes_by_name(fully_migrated_engine, "research_tasks")

        assert "ix_research_tasks_status" in indexes
        assert indexes["ix_research_tasks_status"] == ["status"]

        assert "ix_research_tasks_created_at" in indexes
        assert indexes["ix_research_tasks_created_at"] == ["created_at"]

    def test_creates_research_task_composite_indexes(
        self, fully_migrated_engine
    ):
        """Verify composite indexes on research_tasks with correct column order."""
        indexes = _get_indexes_by_name(fully_migrated_engine, "research_tasks")

        assert "idx_research_task_status_created" in indexes
        assert indexes["idx_research_task_status_created"] == [
            "status",
            "created_at",
        ]

        assert "idx_research_task_priority_status" in indexes
        assert indexes["idx_research_task_priority_status"] == [
            "priority",
            "status",
        ]

    def test_creates_research_history_single_column_indexes(
        self, fully_migrated_engine
    ):
        """Verify single-column indexes on research_history."""
        indexes = _get_indexes_by_name(
            fully_migrated_engine, "research_history"
        )

        assert "ix_research_history_mode" in indexes
        assert indexes["ix_research_history_mode"] == ["mode"]

        assert "ix_research_history_status" in indexes
        assert indexes["ix_research_history_status"] == ["status"]

        assert "ix_research_history_created_at" in indexes
        assert indexes["ix_research_history_created_at"] == ["created_at"]

    def test_creates_research_history_composite_indexes(
        self, fully_migrated_engine
    ):
        """Verify composite indexes on research_history with correct column order."""
        indexes = _get_indexes_by_name(
            fully_migrated_engine, "research_history"
        )

        assert "idx_research_history_status_created" in indexes
        assert indexes["idx_research_history_status_created"] == [
            "status",
            "created_at",
        ]

        assert "idx_research_history_mode_status" in indexes
        assert indexes["idx_research_history_mode_status"] == ["mode", "status"]

    def test_all_indexes_are_non_unique(self, fully_migrated_engine):
        """None of the 9 migration indexes should be unique."""
        insp = inspect(fully_migrated_engine)

        for table_name in ("research_tasks", "research_history"):
            for idx in insp.get_indexes(table_name):
                all_expected = {
                    **RESEARCH_TASKS_INDEXES,
                    **RESEARCH_HISTORY_INDEXES,
                }
                if idx["name"] in all_expected:
                    assert idx["unique"] == 0, (
                        f"Index {idx['name']} should not be unique"
                    )

    def test_total_index_count_research_tasks(self, fully_migrated_engine):
        """research_tasks should have exactly 4 new indexes from this migration."""
        indexes = _get_indexes_by_name(fully_migrated_engine, "research_tasks")
        migration_indexes = {
            name for name in indexes if name in RESEARCH_TASKS_INDEXES
        }
        assert len(migration_indexes) == 4

    def test_total_index_count_research_history(self, fully_migrated_engine):
        """research_history should have exactly 5 new indexes from this migration."""
        indexes = _get_indexes_by_name(
            fully_migrated_engine, "research_history"
        )
        migration_indexes = {
            name for name in indexes if name in RESEARCH_HISTORY_INDEXES
        }
        assert len(migration_indexes) == 5


class TestMigration0003FromFreshDatabase:
    """Tests that verify the full migration chain on a fresh database."""

    def test_fresh_db_full_migration_creates_indexes(self, fresh_engine):
        """Running all migrations on empty DB should create all indexes."""
        run_migrations(fresh_engine)

        for table_name, expected_indexes in [
            ("research_tasks", RESEARCH_TASKS_INDEXES),
            ("research_history", RESEARCH_HISTORY_INDEXES),
        ]:
            indexes = _get_indexes_by_name(fresh_engine, table_name)
            for idx_name, idx_columns in expected_indexes.items():
                assert idx_name in indexes, (
                    f"Missing index {idx_name} on {table_name}"
                )
                assert indexes[idx_name] == idx_columns

    def test_head_revision_is_0004(self):
        """After adding 0004, get_head_revision() should return '0004'."""
        assert get_head_revision() == "0004"

    def test_current_revision_is_0004_after_migrate(self, fresh_engine):
        """After full migration, current revision should be 0004."""
        run_migrations(fresh_engine)
        assert get_current_revision(fresh_engine) == "0004"

    def test_needs_migration_false_after_full_upgrade(self, fresh_engine):
        """After full migration, needs_migration() should return False."""
        run_migrations(fresh_engine)
        assert needs_migration(fresh_engine) is False


class TestMigration0003Idempotency:
    """Tests that verify safe re-runs of migrations."""

    def test_run_migrations_twice_no_error(self, fresh_engine):
        """Calling run_migrations() twice should not raise."""
        run_migrations(fresh_engine)
        run_migrations(fresh_engine)  # Should not raise

    def test_indexes_unchanged_after_double_migration(self, fresh_engine):
        """Indexes should be identical after running migrations twice."""
        run_migrations(fresh_engine)
        indexes_first = {
            "research_tasks": _get_indexes_by_name(
                fresh_engine, "research_tasks"
            ),
            "research_history": _get_indexes_by_name(
                fresh_engine, "research_history"
            ),
        }

        run_migrations(fresh_engine)
        indexes_second = {
            "research_tasks": _get_indexes_by_name(
                fresh_engine, "research_tasks"
            ),
            "research_history": _get_indexes_by_name(
                fresh_engine, "research_history"
            ),
        }

        assert indexes_first == indexes_second

    def test_pre_existing_indexes_not_duplicated(self, migrated_to_0002_engine):
        """Manually creating an index before migration should not cause duplicates."""
        engine = migrated_to_0002_engine

        # Manually create one of the indexes before migration
        with engine.begin() as conn:
            conn.execute(
                text(
                    "CREATE INDEX IF NOT EXISTS ix_research_tasks_status "
                    "ON research_tasks (status)"
                )
            )

        # Now run migration to 0003
        _run_upgrade_to(engine, "0003")

        # Verify no duplicate — only one index with that name
        indexes = _get_indexes_by_name(engine, "research_tasks")
        assert "ix_research_tasks_status" in indexes
        assert indexes["ix_research_tasks_status"] == ["status"]


class TestMigration0003Downgrade:
    """Tests for rollback behavior."""

    def test_downgrade_to_0002_removes_all_research_indexes(
        self, fully_migrated_engine
    ):
        """Downgrade from 0003 to 0002 should remove all 9 indexes."""
        _run_downgrade_to(fully_migrated_engine, "0002")

        # Fresh inspect after DDL
        all_expected = {**RESEARCH_TASKS_INDEXES, **RESEARCH_HISTORY_INDEXES}

        for table_name in ("research_tasks", "research_history"):
            indexes = _get_indexes_by_name(fully_migrated_engine, table_name)
            for idx_name in all_expected:
                assert idx_name not in indexes, (
                    f"Index {idx_name} should have been removed by downgrade"
                )

    def test_downgrade_preserves_tables(self, fully_migrated_engine):
        """Tables should still exist after downgrade (only indexes removed)."""
        _run_downgrade_to(fully_migrated_engine, "0002")

        insp = inspect(fully_migrated_engine)
        assert insp.has_table("research_tasks")
        assert insp.has_table("research_history")

    def test_downgrade_preserves_data(self, fully_migrated_engine):
        """Data inserted before downgrade should survive."""
        engine = fully_migrated_engine

        with engine.begin() as conn:
            conn.execute(
                text(
                    "INSERT INTO research_history (id, query, mode, status, created_at) "
                    "VALUES ('test-downgrade-1', 'test query', 'quick', 'completed', '2025-01-01')"
                )
            )

        _run_downgrade_to(engine, "0002")

        with engine.connect() as conn:
            result = conn.execute(
                text(
                    "SELECT id, query FROM research_history WHERE id = 'test-downgrade-1'"
                )
            ).fetchone()
            assert result is not None
            assert result[0] == "test-downgrade-1"
            assert result[1] == "test query"

    def test_downgrade_then_upgrade_roundtrip(self, fully_migrated_engine):
        """Downgrade to 0002 then upgrade back to 0003 should restore indexes."""
        engine = fully_migrated_engine

        _run_downgrade_to(engine, "0002")
        assert get_current_revision(engine) == "0002"

        _run_upgrade_to(engine, "0003")
        assert get_current_revision(engine) == "0003"

        # Verify indexes are back
        for table_name, expected_indexes in [
            ("research_tasks", RESEARCH_TASKS_INDEXES),
            ("research_history", RESEARCH_HISTORY_INDEXES),
        ]:
            indexes = _get_indexes_by_name(engine, table_name)
            for idx_name in expected_indexes:
                assert idx_name in indexes, (
                    f"Index {idx_name} not restored after roundtrip"
                )


class TestMigration0003DataPreservation:
    """Tests that verify migration is non-destructive."""

    def test_data_preserved_in_research_tasks(self, migrated_to_0002_engine):
        """Data in research_tasks should survive the migration."""
        engine = migrated_to_0002_engine

        with engine.begin() as conn:
            conn.execute(
                text(
                    "INSERT INTO research_tasks (title, status, priority, created_at) "
                    "VALUES ('Test Task', 'pending', 5, '2025-01-01 00:00:00')"
                )
            )

        _run_upgrade_to(engine, "0003")

        with engine.connect() as conn:
            result = conn.execute(
                text("SELECT title, status, priority FROM research_tasks")
            ).fetchone()
            assert result is not None
            assert result[0] == "Test Task"
            assert result[1] == "pending"
            assert result[2] == 5

    def test_data_preserved_in_research_history(self, migrated_to_0002_engine):
        """Data in research_history should survive the migration."""
        engine = migrated_to_0002_engine

        with engine.begin() as conn:
            conn.execute(
                text(
                    "INSERT INTO research_history (id, query, mode, status, created_at) "
                    "VALUES ('preserve-test', 'test query', 'detailed', 'completed', '2025-01-01')"
                )
            )

        _run_upgrade_to(engine, "0003")

        with engine.connect() as conn:
            result = conn.execute(
                text(
                    "SELECT id, query, mode, status FROM research_history "
                    "WHERE id = 'preserve-test'"
                )
            ).fetchone()
            assert result is not None
            assert result[0] == "preserve-test"
            assert result[1] == "test query"
            assert result[2] == "detailed"
            assert result[3] == "completed"

    def test_data_queryable_after_index_creation(self, migrated_to_0002_engine):
        """Queries using indexed columns should work after migration."""
        engine = migrated_to_0002_engine

        with engine.begin() as conn:
            conn.execute(
                text(
                    "INSERT INTO research_tasks (title, status, priority, created_at) "
                    "VALUES ('Task A', 'completed', 10, '2025-01-01 00:00:00')"
                )
            )
            conn.execute(
                text(
                    "INSERT INTO research_tasks (title, status, priority, created_at) "
                    "VALUES ('Task B', 'pending', 5, '2025-01-02 00:00:00')"
                )
            )

        _run_upgrade_to(engine, "0003")

        with engine.connect() as conn:
            # Query using indexed column
            result = conn.execute(
                text(
                    "SELECT title FROM research_tasks WHERE status = 'completed'"
                )
            ).fetchall()
            assert len(result) == 1
            assert result[0][0] == "Task A"

            # Query using composite index columns
            result = conn.execute(
                text(
                    "SELECT title FROM research_tasks "
                    "WHERE priority = 10 AND status = 'completed'"
                )
            ).fetchall()
            assert len(result) == 1
            assert result[0][0] == "Task A"


class TestMigration0003EdgeCases:
    """Tests for edge cases and robustness."""

    def test_partial_tables_only_research_tasks(self, tmp_path):
        """Migration should work when only research_tasks exists."""
        db_path = tmp_path / "partial_tasks_only.db"
        engine = create_engine(f"sqlite:///{db_path}")

        try:
            # Create just research_tasks manually
            with engine.begin() as conn:
                conn.execute(
                    text(
                        "CREATE TABLE research_tasks ("
                        "id INTEGER PRIMARY KEY, "
                        "title VARCHAR(500) NOT NULL, "
                        "status VARCHAR(50), "
                        "priority INTEGER DEFAULT 0, "
                        "created_at DATETIME"
                        ")"
                    )
                )
                # Create alembic_version and stamp at 0002
                conn.execute(
                    text(
                        "CREATE TABLE alembic_version (version_num VARCHAR(32) NOT NULL)"
                    )
                )
                conn.execute(
                    text("INSERT INTO alembic_version VALUES ('0002')")
                )

            _run_upgrade_to(engine, "0003")

            indexes = _get_indexes_by_name(engine, "research_tasks")
            for idx_name in RESEARCH_TASKS_INDEXES:
                assert idx_name in indexes, f"Missing {idx_name}"

            # research_history indexes should not exist (table doesn't exist)
            insp = inspect(engine)
            assert not insp.has_table("research_history")
        finally:
            engine.dispose()

    def test_partial_tables_only_research_history(self, tmp_path):
        """Migration should work when only research_history exists."""
        db_path = tmp_path / "partial_history_only.db"
        engine = create_engine(f"sqlite:///{db_path}")

        try:
            with engine.begin() as conn:
                conn.execute(
                    text(
                        "CREATE TABLE research_history ("
                        "id VARCHAR(36) PRIMARY KEY, "
                        "query TEXT NOT NULL, "
                        "mode TEXT NOT NULL, "
                        "status TEXT NOT NULL, "
                        "created_at TEXT NOT NULL"
                        ")"
                    )
                )
                conn.execute(
                    text(
                        "CREATE TABLE alembic_version (version_num VARCHAR(32) NOT NULL)"
                    )
                )
                conn.execute(
                    text("INSERT INTO alembic_version VALUES ('0002')")
                )

            _run_upgrade_to(engine, "0003")

            indexes = _get_indexes_by_name(engine, "research_history")
            for idx_name in RESEARCH_HISTORY_INDEXES:
                assert idx_name in indexes, f"Missing {idx_name}"

            insp = inspect(engine)
            assert not insp.has_table("research_tasks")
        finally:
            engine.dispose()

    def test_neither_table_exists(self, tmp_path):
        """Migration should be a no-op when neither table exists."""
        db_path = tmp_path / "neither_table.db"
        engine = create_engine(f"sqlite:///{db_path}")

        try:
            with engine.begin() as conn:
                conn.execute(
                    text(
                        "CREATE TABLE alembic_version (version_num VARCHAR(32) NOT NULL)"
                    )
                )
                conn.execute(
                    text("INSERT INTO alembic_version VALUES ('0002')")
                )

            # Should not raise
            _run_upgrade_to(engine, "0003")
            assert get_current_revision(engine) == "0003"
        finally:
            engine.dispose()

    def test_empty_tables_get_indexes(self, migrated_to_0002_engine):
        """Tables with zero rows should still get indexes."""
        engine = migrated_to_0002_engine

        # Verify tables are empty
        with engine.connect() as conn:
            count = conn.execute(
                text("SELECT COUNT(*) FROM research_tasks")
            ).scalar()
            assert count == 0

        _run_upgrade_to(engine, "0003")

        indexes = _get_indexes_by_name(engine, "research_tasks")
        for idx_name in RESEARCH_TASKS_INDEXES:
            assert idx_name in indexes

    def test_in_memory_database(self):
        """Migration should work on in-memory SQLite database."""
        engine = create_engine("sqlite:///:memory:")

        try:
            run_migrations(engine)

            assert get_current_revision(engine) == "0004"

            for table_name, expected_indexes in [
                ("research_tasks", RESEARCH_TASKS_INDEXES),
                ("research_history", RESEARCH_HISTORY_INDEXES),
            ]:
                indexes = _get_indexes_by_name(engine, table_name)
                for idx_name in expected_indexes:
                    assert idx_name in indexes, (
                        f"Missing {idx_name} on {table_name} in memory DB"
                    )
        finally:
            engine.dispose()
