"""End-to-end migration flow tests with index usage verification via EXPLAIN ANALYZE."""

from __future__ import annotations

import pytest
from django.contrib.postgres.operations import (
    AddIndexConcurrently,
    RemoveIndexConcurrently,
)
from django.db import connection, migrations, models
from django.db.migrations.state import ProjectState
from django.db.models import F, Q
from django.db.models.functions import Lower

from paradedb.indexes import BM25Index, IndexExpression

pytestmark = [
    pytest.mark.integration,
    pytest.mark.usefixtures("paradedb_ready"),
]


def _table_exists(table_name: str) -> bool:
    with connection.cursor() as cursor:
        cursor.execute("SELECT to_regclass(%s);", [table_name])
        (regclass,) = cursor.fetchone()
    return regclass == table_name


def _drop_table_if_exists(table_name: str) -> None:
    quoted = connection.ops.quote_name(table_name)
    with connection.cursor() as cursor:
        cursor.execute(f"DROP TABLE IF EXISTS {quoted} CASCADE;")


def _fetch_index_definition(table_name: str, index_name: str) -> str | None:
    with connection.cursor() as cursor:
        cursor.execute(
            """
            SELECT indexdef
            FROM pg_indexes
            WHERE schemaname = current_schema()
              AND tablename = %s
              AND indexname = %s;
            """,
            [table_name, index_name],
        )
        row = cursor.fetchone()
    return row[0] if row else None


def _verify_index_usage(table_name: str, index_name: str) -> bool:  # noqa: ARG001
    """
    Verify BM25 index is actually used by PostgreSQL query planner.

    Returns True if a ParadeDB Custom Scan is found in the execution plan,
    indicating the BM25 index is being utilized for query execution.

    Args:
        table_name: Name of the table to check
        index_name: Name of the BM25 index (kept for API consistency)
    """
    quoted_table = connection.ops.quote_name(table_name)
    with connection.cursor() as cursor:
        cursor.execute(
            f"""
            EXPLAIN (ANALYZE, BUFFERS)
            SELECT id, title, pdb.score(id) as relevance
            FROM {quoted_table}
            WHERE title &&& 'test'
            ORDER BY pdb.score(id) DESC;
            """
        )
        plan_rows = cursor.fetchall()

    plan_text = "\n".join(row[0] for row in plan_rows)

    # ParadeDB Custom Scan node names (from pg_search/src/postgres/customscan/)
    paradedb_scans = (
        "ParadeDB Base Scan",  # BaseScan - standard BM25 queries
        "ParadeDB Aggregate Scan",  # AggregateScan - GROUP BY/aggregates
        "ParadeDB Join Scan",  # JoinScan - join queries with LIMIT
    )
    return any(scan in plan_text for scan in paradedb_scans)


@pytest.mark.django_db(transaction=True)
@pytest.mark.parametrize(
    ("tokenizer", "table_name", "index_name"),
    [
        ("simple", "migtests_items_simple", "migtests_items_simple_bm25_idx"),
        (
            "unicode_words",
            "migtests_items_unicode_words",
            "migtests_items_unicode_words_bm25_idx",
        ),
        ("literal", "migtests_items_literal", "migtests_items_literal_bm25_idx"),
    ],
)
def test_apply_and_unapply_create_model_migration(
    tokenizer: str, table_name: str, index_name: str
) -> None:
    """
    CreateModel with BM25Index migrates forwards/backwards, creates the index,
    and verifies the index is actually used by PostgreSQL query planner.

    Uses transaction=True to allow explicit transaction management and commits,
    simulating real migration behavior where indexes become visible after commit.
    """

    app_label = "migtests"
    # Ensure a clean slate before applying the migration
    _drop_table_if_exists(table_name)
    connection.commit()

    create_model = migrations.CreateModel(
        name=f"MigratedItem_{tokenizer}",
        fields=[
            ("id", models.AutoField(primary_key=True)),
            ("title", models.TextField()),
            ("metadata", models.JSONField(default=dict)),
        ],
        options={
            "db_table": table_name,
            "indexes": [
                BM25Index(
                    fields={
                        "id": {},
                        "title": {"tokenizer": tokenizer},
                        "metadata": {
                            "json_keys": {
                                "color": {"tokenizer": "literal"},
                            }
                        },
                    },
                    key_field="id",
                    name=index_name,
                )
            ],
        },
    )

    from_state = ProjectState()
    to_state = from_state.clone()
    create_model.state_forwards(app_label, to_state)

    forward_applied = False
    with connection.schema_editor(atomic=True) as editor:
        create_model.database_forwards(app_label, editor, from_state, to_state)
        forward_applied = True
    # Schema editor commits on exit

    try:
        # Verify table and index creation after commit
        assert _table_exists(table_name)
        with connection.cursor() as cursor:
            column_names = {
                column.name
                for column in connection.introspection.get_table_description(
                    cursor, table_name
                )
            }
        index_def = _fetch_index_definition(table_name, index_name)
        assert index_def is not None
        assert "USING bm25" in index_def
        normalized_index_def = (
            index_def.replace('"', "")
            .replace("(", "")
            .replace(")", "")
            .replace(" ", "")
        )
        assert f"title::pdb.{tokenizer}" in normalized_index_def, index_def
        assert {"id", "title", "metadata"}.issubset(column_names)

        # Insert test data in a separate transaction (simulates real usage)
        quoted_table = connection.ops.quote_name(table_name)
        with connection.cursor() as cursor:
            cursor.execute(
                f"INSERT INTO {quoted_table} (title, metadata) "
                f"VALUES ('test document', '{{}}'::jsonb), "
                f"('another test', '{{}}'::jsonb), "
                f"('sample text', '{{}}'::jsonb), "
                f"('test', '{{}}'::jsonb);"
            )
        connection.commit()

        # Verify BM25 index is used by query planner in a fresh transaction
        assert _verify_index_usage(table_name, index_name), (
            f"BM25 index {index_name} exists but is not being used by query planner. "
            f"Check EXPLAIN output for Custom Scan (ParadeDBScan)."
        )

    finally:
        if forward_applied and _table_exists(table_name):
            with connection.schema_editor(atomic=True) as editor:
                create_model.database_backwards(app_label, editor, to_state, from_state)
            # Schema editor commits on exit
        else:
            _drop_table_if_exists(table_name)
            connection.commit()
        assert not _table_exists(table_name)


@pytest.mark.django_db(transaction=True)
def test_create_model_migration_with_index_expressions() -> None:
    """CreateModel supports BM25 indexes with computed IndexExpression entries."""
    app_label = "migtests"
    table_name = "migtests_items_index_expression"
    index_name = "migtests_items_index_expression_bm25_idx"

    _drop_table_if_exists(table_name)
    connection.commit()

    create_model = migrations.CreateModel(
        name="MigratedItemIndexExpression",
        fields=[
            ("id", models.AutoField(primary_key=True)),
            ("title", models.TextField()),
            ("rating", models.IntegerField(default=0)),
            ("metadata", models.JSONField(default=dict)),
        ],
        options={
            "db_table": table_name,
            "indexes": [
                BM25Index(
                    fields={
                        "id": {},
                        "title": {"tokenizer": "unicode_words"},
                        "rating": {},
                        "metadata": {},
                    },
                    expressions=[
                        IndexExpression(
                            Lower("title"),
                            alias="title_lower",
                            tokenizer="simple",
                        ),
                        IndexExpression(
                            F("rating") + 1,
                            alias="rating_plus_one",
                        ),
                    ],
                    key_field="id",
                    name=index_name,
                )
            ],
        },
    )

    from_state = ProjectState()
    to_state = from_state.clone()
    create_model.state_forwards(app_label, to_state)

    forward_applied = False
    with connection.schema_editor(atomic=True) as editor:
        create_model.database_forwards(app_label, editor, from_state, to_state)
        forward_applied = True

    try:
        assert _table_exists(table_name)
        index_def = _fetch_index_definition(table_name, index_name)
        assert index_def is not None, "Index was not created"
        normalized_index_def = index_def.lower().replace('"', "").replace(" ", "")
        assert "usingbm25" in normalized_index_def
        assert "lower(title)" in normalized_index_def, index_def
        assert "alias=title_lower" in normalized_index_def, index_def
        assert "rating+1" in normalized_index_def, index_def
        assert "rating_plus_one" in normalized_index_def, index_def
        quoted_table = connection.ops.quote_name(table_name)
        with connection.cursor() as cursor:
            cursor.execute(
                f"INSERT INTO {quoted_table} (title, rating, metadata) VALUES "
                f"('MiXeD Case Title', 4, '{{\"word_count\": 12}}'::jsonb), "
                f"('another record', 2, '{{\"word_count\": 3}}'::jsonb), "
                f"('third row', 9, '{{\"word_count\": 8}}'::jsonb);"
            )
            cursor.execute(
                f"SELECT COUNT(*) FROM {quoted_table} "
                f"WHERE ((LOWER(title))::pdb.alias('title_lower')) &&& 'mixed';"
            )
            (match_count,) = cursor.fetchone()
        connection.commit()

        assert match_count == 1
        assert _verify_index_usage(table_name, index_name)

    finally:
        if forward_applied and _table_exists(table_name):
            with connection.schema_editor(atomic=True) as editor:
                create_model.database_backwards(app_label, editor, to_state, from_state)
        else:
            _drop_table_if_exists(table_name)
            connection.commit()
        assert not _table_exists(table_name)


@pytest.mark.django_db(transaction=True)
def test_multiple_tokenizers_per_field_migration() -> None:
    """A field can be indexed with multiple tokenizers and queried via alias."""
    app_label = "migtests"
    table_name = "migtests_items_multi_tokenizer"
    index_name = "migtests_items_multi_tokenizer_bm25_idx"

    _drop_table_if_exists(table_name)
    connection.commit()

    create_model = migrations.CreateModel(
        name="MigratedItemMultiTokenizer",
        fields=[
            ("id", models.AutoField(primary_key=True)),
            ("title", models.TextField()),
        ],
        options={
            "db_table": table_name,
            "indexes": [
                BM25Index(
                    fields={
                        "id": {},
                        "title": {
                            "tokenizers": [
                                {"tokenizer": "literal"},
                                {"tokenizer": "simple", "alias": "title_simple"},
                            ]
                        },
                    },
                    key_field="id",
                    name=index_name,
                )
            ],
        },
    )

    from_state = ProjectState()
    to_state = from_state.clone()
    create_model.state_forwards(app_label, to_state)

    forward_applied = False
    with connection.schema_editor(atomic=True) as editor:
        create_model.database_forwards(app_label, editor, from_state, to_state)
        forward_applied = True

    try:
        assert _table_exists(table_name)
        index_def = _fetch_index_definition(table_name, index_name)
        assert index_def is not None
        assert "USING bm25" in index_def
        normalized_index_def = (
            index_def.replace('"', "")
            .replace("(", "")
            .replace(")", "")
            .replace(" ", "")
        )
        assert "title::pdb.literal" in normalized_index_def, index_def
        assert "title::pdb.simple" in normalized_index_def, index_def
        assert "alias=title_simple" in normalized_index_def, index_def

        quoted_table = connection.ops.quote_name(table_name)
        with connection.cursor() as cursor:
            cursor.execute(
                f"INSERT INTO {quoted_table} (title) VALUES "
                f"('running shoes'), ('basketball shoes'), ('formal wear');"
            )
            cursor.execute(
                f"SELECT COUNT(*) FROM {quoted_table} WHERE title ||| 'running';"
            )
            (literal_count,) = cursor.fetchone()
            cursor.execute(
                f"SELECT COUNT(*) FROM {quoted_table} "
                f"WHERE (title::pdb.alias('title_simple')) ||| 'running';"
            )
            (simple_alias_count,) = cursor.fetchone()
        connection.commit()

        assert literal_count == 0
        assert simple_alias_count == 1

    finally:
        if forward_applied and _table_exists(table_name):
            with connection.schema_editor(atomic=True) as editor:
                create_model.database_backwards(app_label, editor, to_state, from_state)
        else:
            _drop_table_if_exists(table_name)
            connection.commit()
        assert not _table_exists(table_name)


@pytest.mark.django_db(transaction=True)
def test_multiple_tokenizers_with_ngram_options_migration() -> None:
    """Ngram positional args and named args work in tokenizers DSL."""
    app_label = "migtests"
    table_name = "migtests_items_multi_ngram_tokenizer"
    index_name = "migtests_items_multi_ngram_tokenizer_bm25_idx"

    _drop_table_if_exists(table_name)
    connection.commit()

    create_model = migrations.CreateModel(
        name="MigratedItemMultiNgramTokenizer",
        fields=[
            ("id", models.AutoField(primary_key=True)),
            ("title", models.TextField()),
        ],
        options={
            "db_table": table_name,
            "indexes": [
                BM25Index(
                    fields={
                        "id": {},
                        "title": {
                            "tokenizers": [
                                {"tokenizer": "literal"},
                                {
                                    "tokenizer": "ngram",
                                    "args": [3, 3],
                                    "named_args": {
                                        "prefix_only": True,
                                        "positions": True,
                                    },
                                    "alias": "title_ngram",
                                },
                            ]
                        },
                    },
                    key_field="id",
                    name=index_name,
                )
            ],
        },
    )

    from_state = ProjectState()
    to_state = from_state.clone()
    create_model.state_forwards(app_label, to_state)

    forward_applied = False
    try:
        with connection.schema_editor(atomic=True) as editor:
            create_model.database_forwards(app_label, editor, from_state, to_state)
            forward_applied = True

        assert _table_exists(table_name)
        index_def = _fetch_index_definition(table_name, index_name)
        assert index_def is not None
        normalized_index_def = (
            index_def.replace('"', "")
            .replace("(", "")
            .replace(")", "")
            .replace(" ", "")
        )
        assert "title::pdb.ngram" in normalized_index_def, index_def
        assert "alias=title_ngram" in normalized_index_def, index_def
        assert "prefix_only=true" in normalized_index_def, index_def
        assert "positions=true" in normalized_index_def, index_def

        # Query semantics for ngram aliases are operator-specific and validated
        # in dedicated search-query tests. This migration test verifies DDL only.

    finally:
        if forward_applied and _table_exists(table_name):
            with connection.schema_editor(atomic=True) as editor:
                create_model.database_backwards(app_label, editor, to_state, from_state)
        else:
            _drop_table_if_exists(table_name)
            connection.commit()
        assert not _table_exists(table_name)


@pytest.mark.django_db(transaction=True)
def test_add_index_concurrently() -> None:
    """AddIndexConcurrently creates a BM25 index without holding table locks."""
    app_label = "migtests"
    table_name = "migtests_concurrent"
    index_name = "migtests_concurrent_bm25_idx"

    _drop_table_if_exists(table_name)
    connection.commit()

    bm25_index = BM25Index(
        fields={
            "id": {},
            "title": {"tokenizer": "unicode_words"},
        },
        key_field="id",
        name=index_name,
    )

    # Step 1: Create the table without the index.
    create_model = migrations.CreateModel(
        name="ConcurrentItem",
        fields=[
            ("id", models.AutoField(primary_key=True)),
            ("title", models.TextField()),
        ],
        options={"db_table": table_name},
    )

    from_state = ProjectState()
    to_state = from_state.clone()
    create_model.state_forwards(app_label, to_state)

    with connection.schema_editor(atomic=True) as editor:
        create_model.database_forwards(app_label, editor, from_state, to_state)

    try:
        assert _table_exists(table_name)
        assert _fetch_index_definition(table_name, index_name) is None

        # Step 2: Add the BM25 index concurrently.
        add_index_op = AddIndexConcurrently(
            model_name="concurrentitem",
            index=bm25_index,
        )

        before_state = to_state.clone()
        after_state = before_state.clone()
        add_index_op.state_forwards(app_label, after_state)

        # CONCURRENTLY cannot run inside a transaction.
        with connection.schema_editor(atomic=False) as editor:
            add_index_op.database_forwards(app_label, editor, before_state, after_state)

        index_def = _fetch_index_definition(table_name, index_name)
        assert index_def is not None, "Index was not created"
        assert "USING bm25" in index_def

        # Step 3: Verify the index is functional.
        quoted_table = connection.ops.quote_name(table_name)
        with connection.cursor() as cursor:
            cursor.execute(
                f"INSERT INTO {quoted_table} (title) VALUES "
                f"('test document'), ('another test'), ('sample text');"
            )
        connection.commit()

        assert _verify_index_usage(table_name, index_name)

    finally:
        _drop_table_if_exists(table_name)
        connection.commit()


@pytest.mark.django_db(transaction=True)
def test_remove_index_concurrently() -> None:
    """RemoveIndexConcurrently drops a BM25 index without holding table locks."""
    app_label = "migtests"
    table_name = "migtests_concurrent_rm"
    index_name = "migtests_concurrent_rm_bm25_idx"

    _drop_table_if_exists(table_name)
    connection.commit()

    bm25_index = BM25Index(
        fields={
            "id": {},
            "title": {"tokenizer": "unicode_words"},
        },
        key_field="id",
        name=index_name,
    )

    # Create table with the index.
    create_model = migrations.CreateModel(
        name="ConcurrentRmItem",
        fields=[
            ("id", models.AutoField(primary_key=True)),
            ("title", models.TextField()),
        ],
        options={
            "db_table": table_name,
            "indexes": [bm25_index],
        },
    )

    from_state = ProjectState()
    to_state = from_state.clone()
    create_model.state_forwards(app_label, to_state)

    with connection.schema_editor(atomic=True) as editor:
        create_model.database_forwards(app_label, editor, from_state, to_state)

    try:
        assert _fetch_index_definition(table_name, index_name) is not None

        # Remove the index concurrently.
        remove_op = RemoveIndexConcurrently(
            model_name="concurrentrmitem",
            name=index_name,
        )

        before_state = to_state.clone()
        after_state = before_state.clone()
        remove_op.state_forwards(app_label, after_state)

        with connection.schema_editor(atomic=False) as editor:
            remove_op.database_forwards(app_label, editor, before_state, after_state)

        assert _fetch_index_definition(table_name, index_name) is None, (
            "Index was not dropped"
        )

    finally:
        _drop_table_if_exists(table_name)
        connection.commit()


@pytest.mark.django_db(transaction=True)
def test_add_partial_index() -> None:
    """BM25Index with condition creates a partial index with a WHERE clause."""
    app_label = "migtests"
    table_name = "migtests_partial"
    index_name = "migtests_partial_bm25_idx"

    _drop_table_if_exists(table_name)
    connection.commit()

    bm25_index = BM25Index(
        fields={
            "id": {},
            "title": {"tokenizer": "unicode_words"},
        },
        key_field="id",
        name=index_name,
        condition=Q(is_active=True),
    )

    create_model = migrations.CreateModel(
        name="PartialItem",
        fields=[
            ("id", models.AutoField(primary_key=True)),
            ("title", models.TextField()),
            ("is_active", models.BooleanField(default=True)),
        ],
        options={
            "db_table": table_name,
            "indexes": [bm25_index],
        },
    )

    from_state = ProjectState()
    to_state = from_state.clone()
    create_model.state_forwards(app_label, to_state)

    forward_applied = False
    with connection.schema_editor(atomic=True) as editor:
        create_model.database_forwards(app_label, editor, from_state, to_state)
        forward_applied = True

    try:
        assert _table_exists(table_name)
        index_def = _fetch_index_definition(table_name, index_name)
        assert index_def is not None, "Index was not created"
        assert "USING bm25" in index_def
        assert "WHERE" in index_def

        # Insert test data and verify the index is functional
        quoted_table = connection.ops.quote_name(table_name)
        with connection.cursor() as cursor:
            cursor.execute(
                f"INSERT INTO {quoted_table} (title, is_active) VALUES "
                f"('test document', true), ('inactive doc', false), "
                f"('another test', true);"
            )
        connection.commit()

        # Verify the partial index is queryable
        with connection.cursor() as cursor:
            cursor.execute(
                f"SELECT COUNT(*) FROM {quoted_table} WHERE title &&& 'test';"
            )
            (count,) = cursor.fetchone()
        # Only active rows are indexed; 'test document' and 'another test'
        assert count == 2

    finally:
        if forward_applied and _table_exists(table_name):
            with connection.schema_editor(atomic=True) as editor:
                create_model.database_backwards(app_label, editor, to_state, from_state)
        else:
            _drop_table_if_exists(table_name)
            connection.commit()
        assert not _table_exists(table_name)
