from typing import Any, TYPE_CHECKING

from django.contrib.postgres.indexes import (
    GinIndex,
    GistIndex,
    BTreeIndex,  # noqa  convenience import
)
from django.db.models import Index, Model
from django.db.backends.ddl_references import Statement
from django.db.backends.base.schema import BaseDatabaseSchemaEditor

if TYPE_CHECKING:
    IndexBase = Index
else:
    IndexBase = object


class MultiColumnSql(IndexBase):
    """
    Index normal column and UPPER(column_name) so the index can be used
    for both <fieldname>__trigram_similar and <fieldname>__icontains queries.
    """

    def create_sql(
        self,
        model: type[Model],
        schema_editor: BaseDatabaseSchemaEditor,
        using: str = '',
        **kwargs: Any,
    ) -> Statement:
        """
        For each column in the index add the same column with UPPER(column)
        """
        statement = super().create_sql(model, schema_editor, using=using, **kwargs)  # noqa

        # get columns, which should be comma-separated column name & ops class, per index type.
        # "name" gist_trgm_ops , "code" gist_trgm_ops
        # "name" gin_trgm_ops , "code" gin_trgm_ops
        standard_columns = str(statement.parts["columns"])

        upper_columns = []
        for column in standard_columns.split(','):
            name, op_class = column.strip().split()
            upper_columns.append(f'UPPER({name}) {op_class}')

        statement.parts["columns_upper"] = f'{standard_columns} , {" , ".join(upper_columns)}'  # type: ignore[assignment]
        statement.template = statement.template.replace('columns', 'columns_upper')

        return statement


class UpperGistIndex(MultiColumnSql, GistIndex):
    """
    Index normal column and UPPER(column_name) so the index can be used
    for both <fieldname>__trigram_similar and <fieldname>__icontains queries.
    """

    def __init__(
        self,
        *,
        buffering: bool | None = None,
        fillfactor: int | None = None,
        **kwargs: Any,
    ):  # noqa
        super().__init__(**kwargs)


class UpperGinIndex(MultiColumnSql, GinIndex):
    pass
