"""Add search vectors

Revision ID: 598f0b91204a
Revises: cf31af43f14a
Create Date: 2025-12-10 17:37:18.771934

"""

import sqlalchemy as sa
from alembic import op
from alembic_utils.pg_function import PGFunction
from alembic_utils.pg_trigger import PGTrigger
from sqlalchemy import text as sql_text
from sqlalchemy.dialects import postgresql

# Polar Custom Imports

# revision identifiers, used by Alembic.
revision = "598f0b91204a"
down_revision = "cf31af43f14a"
branch_labels: tuple[str] | None = None
depends_on: tuple[str] | None = None


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column(
        "customers", sa.Column("search_vector", postgresql.TSVECTOR(), nullable=True)
    )
    op.add_column(
        "orders", sa.Column("search_vector", postgresql.TSVECTOR(), nullable=True)
    )
    op.add_column(
        "products", sa.Column("search_vector", postgresql.TSVECTOR(), nullable=True)
    )
    public_customers_search_vector_update = PGFunction(
        schema="public",
        signature="customers_search_vector_update()",
        definition="RETURNS trigger AS $$\n    BEGIN\n        NEW.search_vector := to_tsvector('simple', coalesce(NEW.name, ''));\n        RETURN NEW;\n    END\n    $$ LANGUAGE plpgsql",
    )
    op.create_entity(public_customers_search_vector_update)

    public_products_search_vector_update = PGFunction(
        schema="public",
        signature="products_search_vector_update()",
        definition="RETURNS trigger AS $$\n    BEGIN\n        NEW.search_vector := to_tsvector('english', coalesce(NEW.name, '') || ' ' || coalesce(NEW.description, ''));\n        RETURN NEW;\n    END\n    $$ LANGUAGE plpgsql",
    )
    op.create_entity(public_products_search_vector_update)

    public_orders_search_vector_update = PGFunction(
        schema="public",
        signature="orders_search_vector_update()",
        definition="RETURNS trigger AS $$\n    BEGIN\n        NEW.search_vector := to_tsvector('simple', coalesce(NEW.invoice_number, '') || ' ' || coalesce(NEW.stripe_invoice_id, ''));\n        RETURN NEW;\n    END\n    $$ LANGUAGE plpgsql",
    )
    op.create_entity(public_orders_search_vector_update)

    public_customers_customers_search_vector_trigger = PGTrigger(
        schema="public",
        signature="customers_search_vector_trigger",
        on_entity="public.customers",
        is_constraint=False,
        definition="BEFORE INSERT OR UPDATE ON customers\n    FOR EACH ROW EXECUTE FUNCTION customers_search_vector_update()",
    )
    op.create_entity(public_customers_customers_search_vector_trigger)

    public_products_products_search_vector_trigger = PGTrigger(
        schema="public",
        signature="products_search_vector_trigger",
        on_entity="public.products",
        is_constraint=False,
        definition="BEFORE INSERT OR UPDATE ON products\n    FOR EACH ROW EXECUTE FUNCTION products_search_vector_update()",
    )
    op.create_entity(public_products_products_search_vector_trigger)

    public_orders_orders_search_vector_trigger = PGTrigger(
        schema="public",
        signature="orders_search_vector_trigger",
        on_entity="public.orders",
        is_constraint=False,
        definition="BEFORE INSERT OR UPDATE ON orders\n    FOR EACH ROW EXECUTE FUNCTION orders_search_vector_update()",
    )
    op.create_entity(public_orders_orders_search_vector_trigger)

    with op.get_context().autocommit_block():
        op.create_index(
            "ix_customers_search_vector",
            "customers",
            ["search_vector"],
            unique=False,
            postgresql_using="gin",
            postgresql_concurrently=True,
        )
        op.create_index(
            "ix_orders_search_vector",
            "orders",
            ["search_vector"],
            unique=False,
            postgresql_using="gin",
            postgresql_concurrently=True,
        )
        op.create_index(
            "ix_products_search_vector",
            "products",
            ["search_vector"],
            unique=False,
            postgresql_using="gin",
            postgresql_concurrently=True,
        )

    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    public_orders_orders_search_vector_trigger = PGTrigger(
        schema="public",
        signature="orders_search_vector_trigger",
        on_entity="public.orders",
        is_constraint=False,
        definition="BEFORE INSERT OR UPDATE ON orders\n    FOR EACH ROW EXECUTE FUNCTION orders_search_vector_update()",
    )
    op.drop_entity(public_orders_orders_search_vector_trigger)

    public_products_products_search_vector_trigger = PGTrigger(
        schema="public",
        signature="products_search_vector_trigger",
        on_entity="public.products",
        is_constraint=False,
        definition="BEFORE INSERT OR UPDATE ON products\n    FOR EACH ROW EXECUTE FUNCTION products_search_vector_update()",
    )
    op.drop_entity(public_products_products_search_vector_trigger)

    public_customers_customers_search_vector_trigger = PGTrigger(
        schema="public",
        signature="customers_search_vector_trigger",
        on_entity="public.customers",
        is_constraint=False,
        definition="BEFORE INSERT OR UPDATE ON customers\n    FOR EACH ROW EXECUTE FUNCTION customers_search_vector_update()",
    )
    op.drop_entity(public_customers_customers_search_vector_trigger)

    public_orders_search_vector_update = PGFunction(
        schema="public",
        signature="orders_search_vector_update()",
        definition="RETURNS trigger AS $$\n    BEGIN\n        NEW.search_vector := to_tsvector('simple', coalesce(NEW.invoice_number, '') || ' ' || coalesce(NEW.stripe_invoice_id, ''));\n        RETURN NEW;\n    END\n    $$ LANGUAGE plpgsql",
    )
    op.drop_entity(public_orders_search_vector_update)

    public_products_search_vector_update = PGFunction(
        schema="public",
        signature="products_search_vector_update()",
        definition="RETURNS trigger AS $$\n    BEGIN\n        NEW.search_vector := to_tsvector('english', coalesce(NEW.name, '') || ' ' || coalesce(NEW.description, ''));\n        RETURN NEW;\n    END\n    $$ LANGUAGE plpgsql",
    )
    op.drop_entity(public_products_search_vector_update)

    public_customers_search_vector_update = PGFunction(
        schema="public",
        signature="customers_search_vector_update()",
        definition="RETURNS trigger AS $$\n    BEGIN\n        NEW.search_vector := to_tsvector('simple', coalesce(NEW.name, ''));\n        RETURN NEW;\n    END\n    $$ LANGUAGE plpgsql",
    )
    op.drop_entity(public_customers_search_vector_update)

    op.drop_index(
        "ix_products_search_vector", table_name="products", postgresql_using="gin"
    )
    op.drop_column("products", "search_vector")
    op.drop_index(
        "ix_orders_search_vector", table_name="orders", postgresql_using="gin"
    )
    op.drop_column("orders", "search_vector")
    op.drop_index(
        "ix_customers_search_vector", table_name="customers", postgresql_using="gin"
    )
    op.drop_column("customers", "search_vector")
    # ### end Alembic commands ###
