"""message search index outside the hypertable, versions searchable

Revision ID: d8f1a6c3e5b7
Revises: c9e2a7f4b1d8
Create Date: 2026-10-19 03:00:00.000000

"""

from collections.abc import Sequence

from alembic import op

revision: str = "d8f1a6c3e5b7"
down_revision: str | None = "c9e2a7f4b1d8"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None

_VECTOR = """
CREATE FUNCTION search_vector(body text) RETURNS tsvector
LANGUAGE sql IMMUTABLE PARALLEL SAFE AS $$
    SELECT to_tsvector('russian', translate(coalesce(body, ''), 'Ёё', 'Ее'))
        || to_tsvector('simple', translate(coalesce(body, ''), 'Ёё', 'Ее'))
$$
"""

_TABLE = """
CREATE TABLE message_search (
    account_id integer NOT NULL,
    chat_id bigint NOT NULL,
    message_id bigint NOT NULL,
    date timestamptz NOT NULL,
    sender_id bigint,
    text text,
    has_media boolean NOT NULL DEFAULT false,
    edited_at timestamptz,
    deleted_at timestamptz,
    tsv tsvector GENERATED ALWAYS AS (search_vector(text)) STORED,
    PRIMARY KEY (account_id, chat_id, message_id, date)
)
"""

_BACKFILL = """
INSERT INTO message_search (account_id, chat_id, message_id, date, sender_id,
                            text, has_media, edited_at, deleted_at)
SELECT account_id, chat_id, message_id, date, sender_id,
       text, has_media, edited_at, deleted_at
FROM messages
"""

_SYNC = """
CREATE FUNCTION message_search_sync() RETURNS trigger LANGUAGE plpgsql AS $$
BEGIN
    IF TG_OP = 'DELETE' THEN
        DELETE FROM message_search
        WHERE account_id = OLD.account_id AND chat_id = OLD.chat_id
          AND message_id = OLD.message_id AND date = OLD.date;
        RETURN NULL;
    END IF;
    INSERT INTO message_search (account_id, chat_id, message_id, date,
                                sender_id, text, has_media, edited_at,
                                deleted_at)
    VALUES (NEW.account_id, NEW.chat_id, NEW.message_id, NEW.date,
            NEW.sender_id, NEW.text, NEW.has_media, NEW.edited_at,
            NEW.deleted_at)
    ON CONFLICT (account_id, chat_id, message_id, date) DO UPDATE SET
        sender_id = EXCLUDED.sender_id,
        text = EXCLUDED.text,
        has_media = EXCLUDED.has_media,
        edited_at = EXCLUDED.edited_at,
        deleted_at = EXCLUDED.deleted_at;
    RETURN NULL;
END;
$$
"""

_CHANGED = (
    "OLD.text IS DISTINCT FROM NEW.text "
    "OR OLD.sender_id IS DISTINCT FROM NEW.sender_id "
    "OR OLD.has_media IS DISTINCT FROM NEW.has_media "
    "OR OLD.edited_at IS DISTINCT FROM NEW.edited_at "
    "OR OLD.deleted_at IS DISTINCT FROM NEW.deleted_at"
)


def upgrade() -> None:
    op.execute(_VECTOR)
    op.execute(_TABLE)
    op.execute(_BACKFILL)
    op.execute("CREATE INDEX ix_message_search_tsv ON message_search USING gin (tsv)")
    op.execute(
        "CREATE INDEX ix_message_search_recent ON message_search "
        "(account_id, date DESC, chat_id DESC, message_id DESC)"
    )
    op.execute(
        "CREATE INDEX ix_message_search_sender ON message_search "
        "(account_id, sender_id, date DESC)"
    )
    op.execute(_SYNC)
    op.execute(
        "CREATE TRIGGER messages_search_insert AFTER INSERT ON messages "
        "FOR EACH ROW EXECUTE FUNCTION message_search_sync()"
    )
    op.execute(
        "CREATE TRIGGER messages_search_update AFTER UPDATE ON messages "
        f"FOR EACH ROW WHEN ({_CHANGED}) EXECUTE FUNCTION message_search_sync()"
    )
    op.execute(
        "CREATE TRIGGER messages_search_delete AFTER DELETE ON messages "
        "FOR EACH ROW EXECUTE FUNCTION message_search_sync()"
    )
    op.execute(
        "CREATE INDEX ix_message_versions_tsv ON message_versions USING gin "
        "(search_vector(text))"
    )
    op.execute("DROP INDEX IF EXISTS ix_media_tsv")
    op.execute(
        "CREATE INDEX ix_media_tsv ON media USING gin (search_vector(extracted_text))"
    )
    op.execute("DROP INDEX IF EXISTS ix_messages_tsv")


def downgrade() -> None:
    op.execute(
        "CREATE INDEX ix_messages_tsv ON messages USING gin "
        "(to_tsvector('russian', coalesce(text, '')))"
    )
    op.execute("DROP INDEX IF EXISTS ix_media_tsv")
    op.execute(
        "CREATE INDEX ix_media_tsv ON media USING gin "
        "(to_tsvector('russian', coalesce(extracted_text, '')))"
    )
    op.execute("DROP INDEX IF EXISTS ix_message_versions_tsv")
    op.execute("DROP TRIGGER IF EXISTS messages_search_delete ON messages")
    op.execute("DROP TRIGGER IF EXISTS messages_search_update ON messages")
    op.execute("DROP TRIGGER IF EXISTS messages_search_insert ON messages")
    op.execute("DROP FUNCTION IF EXISTS message_search_sync()")
    op.execute("DROP TABLE IF EXISTS message_search")
    op.execute("DROP FUNCTION IF EXISTS search_vector(text)")