"""Add Feishu identities, personalization, and durable subscriptions.

Revision ID: 202607260003
Revises: 202607260002
Create Date: 2026-07-26
"""

from alembic import op
import sqlalchemy as sa


revision = "202607260003"
down_revision = "202607260002"
branch_labels = None
depends_on = None

AI_MEMORY_TABLE = "ai_memory_entries"
MARKET_WATCHLIST_TABLE = "market_watchlists"


def _create_indexes(
    table_name: str,
    definitions: tuple[tuple[str, tuple[str, ...], bool], ...],
) -> None:
    for name, columns, unique in definitions:
        op.create_index(op.f(name), table_name, list(columns), unique=unique)


def upgrade() -> None:
    _create_feishu_users()
    _upgrade_ai_memory()
    _upgrade_market_watchlists()
    _create_personalization_tables()
    _create_subscription_tables()


def _create_feishu_users() -> None:
    op.create_table(
        "feishu_users",
        sa.Column("id", sa.Integer(), autoincrement=True, nullable=False),
        sa.Column("code", sa.String(length=64), nullable=False),
        sa.Column("tenant_key", sa.String(length=128), nullable=False),
        sa.Column("open_id", sa.String(length=128), nullable=False),
        sa.Column("union_id", sa.String(length=128), nullable=True),
        sa.Column("user_id", sa.String(length=128), nullable=True),
        sa.Column("role", sa.String(length=32), nullable=False),
        sa.Column("status", sa.String(length=32), nullable=False),
        sa.Column("timezone", sa.String(length=64), nullable=False),
        sa.Column("quiet_hours_start", sa.Time(), nullable=True),
        sa.Column("quiet_hours_end", sa.Time(), nullable=True),
        sa.Column("last_active_at", sa.DateTime(), nullable=False),
        sa.Column("created_at", sa.DateTime(), nullable=False),
        sa.Column("updated_at", sa.DateTime(), nullable=False),
        sa.PrimaryKeyConstraint("id"),
        sa.UniqueConstraint(
            "tenant_key",
            "open_id",
            name="uq_feishu_user_tenant_open_id",
        ),
    )
    _create_indexes(
        "feishu_users",
        (
            ("ix_feishu_users_code", ("code",), True),
            ("ix_feishu_users_created_at", ("created_at",), False),
            ("ix_feishu_users_last_active_at", ("last_active_at",), False),
            ("ix_feishu_users_open_id", ("open_id",), False),
            ("ix_feishu_users_role", ("role",), False),
            ("ix_feishu_users_status", ("status",), False),
            ("ix_feishu_users_tenant_key", ("tenant_key",), False),
            ("ix_feishu_users_union_id", ("union_id",), False),
            ("ix_feishu_users_user_id", ("user_id",), False),
        ),
    )


def _upgrade_ai_memory() -> None:
    op.drop_index(
        "ix_ai_memory_entries_fingerprint",
        table_name=AI_MEMORY_TABLE,
    )
    with op.batch_alter_table(AI_MEMORY_TABLE) as batch_op:
        batch_op.add_column(
            sa.Column("owner_id", sa.Integer(), nullable=True)
        )
        batch_op.add_column(
            sa.Column(
                "kind",
                sa.String(length=32),
                nullable=False,
                server_default="memory",
            )
        )
        batch_op.create_foreign_key(
            "fk_ai_memory_entries_owner_id",
            "feishu_users",
            ["owner_id"],
            ["id"],
            ondelete="CASCADE",
        )
        batch_op.create_unique_constraint(
            "uq_ai_memory_owner_fingerprint",
            ["owner_id", "fingerprint"],
        )
        batch_op.create_index(
            "ix_ai_memory_entries_fingerprint",
            ["fingerprint"],
            unique=False,
        )
        batch_op.create_index(
            "ix_ai_memory_entries_owner_id",
            ["owner_id"],
            unique=False,
        )
        batch_op.create_index(
            "ix_ai_memory_entries_kind",
            ["kind"],
            unique=False,
        )

    connection = op.get_bind()
    connection.execute(
        sa.text(
            """
            UPDATE ai_memory_entries
            SET
                owner_id = NULL,
                kind = 'company_rule',
                source = 'legacy_company',
                status = 'active'
            WHERE source = 'user_rule'
            """
        )
    )
    connection.execute(
        sa.text(
            """
            UPDATE ai_memory_entries
            SET kind = 'memory', status = 'archived'
            WHERE owner_id IS NULL AND source IN ('auto', 'hermes')
            """
        )
    )
    with op.batch_alter_table(AI_MEMORY_TABLE) as batch_op:
        batch_op.alter_column(
            "kind",
            existing_type=sa.String(length=32),
            nullable=False,
            server_default=None,
        )


def _upgrade_market_watchlists() -> None:
    with op.batch_alter_table(MARKET_WATCHLIST_TABLE) as batch_op:
        batch_op.drop_constraint(
            "uq_market_watchlist_actor_symbol",
            type_="unique",
        )
        batch_op.add_column(sa.Column("owner_id", sa.Integer(), nullable=True))
        batch_op.create_foreign_key(
            "fk_market_watchlists_owner_id",
            "feishu_users",
            ["owner_id"],
            ["id"],
            ondelete="CASCADE",
        )
        batch_op.create_unique_constraint(
            "uq_market_watchlist_owner_symbol",
            ["owner_id", "symbol"],
        )
        batch_op.create_index(
            "ix_market_watchlists_owner_id",
            ["owner_id"],
            unique=False,
        )


def _create_personalization_tables() -> None:
    op.create_table(
        "user_preferences",
        sa.Column("id", sa.Integer(), autoincrement=True, nullable=False),
        sa.Column("code", sa.String(length=64), nullable=False),
        sa.Column("owner_id", sa.Integer(), nullable=False),
        sa.Column("category", sa.String(length=32), nullable=False),
        sa.Column("value", sa.Text(), nullable=False),
        sa.Column("normalized_value", sa.String(length=1000), nullable=False),
        sa.Column("source", sa.String(length=32), nullable=False),
        sa.Column("created_at", sa.DateTime(), nullable=False),
        sa.Column("updated_at", sa.DateTime(), nullable=False),
        sa.ForeignKeyConstraint(
            ["owner_id"],
            ["feishu_users.id"],
            name="fk_user_preferences_owner_id",
            ondelete="CASCADE",
        ),
        sa.PrimaryKeyConstraint("id"),
        sa.UniqueConstraint(
            "owner_id",
            "category",
            "normalized_value",
            name="uq_user_preference_owner_category_value",
        ),
    )
    _create_indexes(
        "user_preferences",
        (
            ("ix_user_preferences_category", ("category",), False),
            ("ix_user_preferences_code", ("code",), True),
            ("ix_user_preferences_created_at", ("created_at",), False),
            ("ix_user_preferences_owner_id", ("owner_id",), False),
            ("ix_user_preferences_source", ("source",), False),
        ),
    )

    op.create_table(
        "ai_conversations",
        sa.Column("id", sa.Integer(), autoincrement=True, nullable=False),
        sa.Column("code", sa.String(length=64), nullable=False),
        sa.Column("owner_id", sa.Integer(), nullable=False),
        sa.Column("chat_type", sa.String(length=32), nullable=False),
        sa.Column("chat_key", sa.String(length=256), nullable=False),
        sa.Column("created_at", sa.DateTime(), nullable=False),
        sa.Column("updated_at", sa.DateTime(), nullable=False),
        sa.ForeignKeyConstraint(
            ["owner_id"],
            ["feishu_users.id"],
            name="fk_ai_conversations_owner_id",
            ondelete="CASCADE",
        ),
        sa.PrimaryKeyConstraint("id"),
        sa.UniqueConstraint(
            "owner_id",
            "chat_type",
            "chat_key",
            name="uq_ai_conversation_owner_chat",
        ),
    )
    _create_indexes(
        "ai_conversations",
        (
            ("ix_ai_conversations_chat_key", ("chat_key",), False),
            ("ix_ai_conversations_chat_type", ("chat_type",), False),
            ("ix_ai_conversations_code", ("code",), True),
            ("ix_ai_conversations_created_at", ("created_at",), False),
            ("ix_ai_conversations_owner_id", ("owner_id",), False),
            ("ix_ai_conversations_updated_at", ("updated_at",), False),
        ),
    )

    op.create_table(
        "ai_conversation_messages",
        sa.Column("id", sa.Integer(), autoincrement=True, nullable=False),
        sa.Column("conversation_id", sa.Integer(), nullable=False),
        sa.Column("role", sa.String(length=32), nullable=False),
        sa.Column("content", sa.Text(), nullable=False),
        sa.Column("created_at", sa.DateTime(), nullable=False),
        sa.ForeignKeyConstraint(
            ["conversation_id"],
            ["ai_conversations.id"],
            name="fk_ai_conversation_messages_conversation_id",
            ondelete="CASCADE",
        ),
        sa.PrimaryKeyConstraint("id"),
    )
    _create_indexes(
        "ai_conversation_messages",
        (
            (
                "ix_ai_conversation_messages_conversation_id",
                ("conversation_id",),
                False,
            ),
            ("ix_ai_conversation_messages_created_at", ("created_at",), False),
            ("ix_ai_conversation_messages_role", ("role",), False),
        ),
    )

    op.create_table(
        "personal_data_erasure_requests",
        sa.Column("id", sa.Integer(), autoincrement=True, nullable=False),
        sa.Column("owner_id", sa.Integer(), nullable=False),
        sa.Column("token_hash", sa.String(length=64), nullable=False),
        sa.Column("expires_at", sa.DateTime(), nullable=False),
        sa.Column("created_at", sa.DateTime(), nullable=False),
        sa.Column("updated_at", sa.DateTime(), nullable=False),
        sa.ForeignKeyConstraint(
            ["owner_id"],
            ["feishu_users.id"],
            name="fk_personal_data_erasure_requests_owner_id",
            ondelete="CASCADE",
        ),
        sa.PrimaryKeyConstraint("id"),
    )
    _create_indexes(
        "personal_data_erasure_requests",
        (
            (
                "ix_personal_data_erasure_requests_created_at",
                ("created_at",),
                False,
            ),
            (
                "ix_personal_data_erasure_requests_expires_at",
                ("expires_at",),
                False,
            ),
            (
                "ix_personal_data_erasure_requests_owner_id",
                ("owner_id",),
                True,
            ),
        ),
    )


def _create_subscription_tables() -> None:
    op.create_table(
        "push_subscriptions",
        sa.Column("id", sa.Integer(), autoincrement=True, nullable=False),
        sa.Column("code", sa.String(length=64), nullable=False),
        sa.Column("owner_id", sa.Integer(), nullable=False),
        sa.Column("target_type", sa.String(length=16), nullable=False),
        sa.Column("target_id", sa.String(length=256), nullable=False),
        sa.Column("prompt", sa.Text(), nullable=False),
        sa.Column("schedule_type", sa.String(length=32), nullable=False),
        sa.Column("schedule_config", sa.JSON(), nullable=False),
        sa.Column("timezone", sa.String(length=64), nullable=False),
        sa.Column("next_run_at", sa.DateTime(), nullable=True),
        sa.Column("status", sa.String(length=32), nullable=False),
        sa.Column("consented_at", sa.DateTime(), nullable=False),
        sa.Column("last_run_at", sa.DateTime(), nullable=True),
        sa.Column("locked_by", sa.String(length=128), nullable=True),
        sa.Column("locked_until", sa.DateTime(), nullable=True),
        sa.Column("created_at", sa.DateTime(), nullable=False),
        sa.Column("updated_at", sa.DateTime(), nullable=False),
        sa.ForeignKeyConstraint(
            ["owner_id"],
            ["feishu_users.id"],
            name="fk_push_subscriptions_owner_id",
            ondelete="CASCADE",
        ),
        sa.PrimaryKeyConstraint("id"),
    )
    _create_indexes(
        "push_subscriptions",
        (
            ("ix_push_subscriptions_code", ("code",), True),
            ("ix_push_subscriptions_created_at", ("created_at",), False),
            ("ix_push_subscriptions_locked_by", ("locked_by",), False),
            ("ix_push_subscriptions_locked_until", ("locked_until",), False),
            ("ix_push_subscriptions_next_run_at", ("next_run_at",), False),
            ("ix_push_subscriptions_owner_id", ("owner_id",), False),
            ("ix_push_subscriptions_schedule_type", ("schedule_type",), False),
            ("ix_push_subscriptions_status", ("status",), False),
            ("ix_push_subscriptions_target_type", ("target_type",), False),
        ),
    )

    op.create_table(
        "push_deliveries",
        sa.Column("id", sa.Integer(), autoincrement=True, nullable=False),
        sa.Column("code", sa.String(length=64), nullable=False),
        sa.Column("subscription_id", sa.Integer(), nullable=False),
        sa.Column("scheduled_for", sa.DateTime(), nullable=False),
        sa.Column("idempotency_key", sa.String(length=64), nullable=False),
        sa.Column("message_uuid", sa.String(length=36), nullable=False),
        sa.Column("status", sa.String(length=32), nullable=False),
        sa.Column("attempt_count", sa.Integer(), nullable=False),
        sa.Column("next_attempt_at", sa.DateTime(), nullable=True),
        sa.Column("rendered_content", sa.Text(), nullable=True),
        sa.Column("provider_message_id", sa.String(length=256), nullable=True),
        sa.Column("last_error", sa.Text(), nullable=True),
        sa.Column("locked_by", sa.String(length=128), nullable=True),
        sa.Column("locked_until", sa.DateTime(), nullable=True),
        sa.Column("sent_at", sa.DateTime(), nullable=True),
        sa.Column("created_at", sa.DateTime(), nullable=False),
        sa.Column("updated_at", sa.DateTime(), nullable=False),
        sa.ForeignKeyConstraint(
            ["subscription_id"],
            ["push_subscriptions.id"],
            name="fk_push_deliveries_subscription_id",
            ondelete="CASCADE",
        ),
        sa.PrimaryKeyConstraint("id"),
        sa.UniqueConstraint(
            "subscription_id",
            "scheduled_for",
            name="uq_push_delivery_subscription_schedule",
        ),
    )
    _create_indexes(
        "push_deliveries",
        (
            ("ix_push_deliveries_code", ("code",), True),
            ("ix_push_deliveries_created_at", ("created_at",), False),
            (
                "ix_push_deliveries_idempotency_key",
                ("idempotency_key",),
                True,
            ),
            ("ix_push_deliveries_locked_by", ("locked_by",), False),
            ("ix_push_deliveries_locked_until", ("locked_until",), False),
            ("ix_push_deliveries_message_uuid", ("message_uuid",), True),
            (
                "ix_push_deliveries_next_attempt_at",
                ("next_attempt_at",),
                False,
            ),
            (
                "ix_push_deliveries_provider_message_id",
                ("provider_message_id",),
                False,
            ),
            ("ix_push_deliveries_scheduled_for", ("scheduled_for",), False),
            ("ix_push_deliveries_sent_at", ("sent_at",), False),
            ("ix_push_deliveries_status", ("status",), False),
            (
                "ix_push_deliveries_subscription_id",
                ("subscription_id",),
                False,
            ),
        ),
    )


def downgrade() -> None:
    op.drop_table("push_deliveries")
    op.drop_table("push_subscriptions")
    op.drop_table("personal_data_erasure_requests")
    op.drop_table("ai_conversation_messages")
    op.drop_table("ai_conversations")
    op.drop_table("user_preferences")

    with op.batch_alter_table(MARKET_WATCHLIST_TABLE) as batch_op:
        batch_op.drop_constraint(
            "uq_market_watchlist_owner_symbol",
            type_="unique",
        )
        batch_op.drop_constraint(
            "fk_market_watchlists_owner_id",
            type_="foreignkey",
        )
        batch_op.drop_index("ix_market_watchlists_owner_id")
        batch_op.drop_column("owner_id")
        batch_op.create_unique_constraint(
            "uq_market_watchlist_actor_symbol",
            ["actor", "symbol"],
        )

    op.execute(
        sa.text(
            """
            UPDATE ai_memory_entries
            SET source = 'user_rule'
            WHERE owner_id IS NULL
              AND kind = 'company_rule'
              AND source = 'legacy_company'
            """
        )
    )
    with op.batch_alter_table(AI_MEMORY_TABLE) as batch_op:
        batch_op.drop_constraint(
            "uq_ai_memory_owner_fingerprint",
            type_="unique",
        )
        batch_op.drop_constraint(
            "fk_ai_memory_entries_owner_id",
            type_="foreignkey",
        )
        batch_op.drop_index("ix_ai_memory_entries_fingerprint")
        batch_op.drop_index("ix_ai_memory_entries_kind")
        batch_op.drop_index("ix_ai_memory_entries_owner_id")
        batch_op.drop_column("kind")
        batch_op.drop_column("owner_id")
        batch_op.create_index(
            "ix_ai_memory_entries_fingerprint",
            ["fingerprint"],
            unique=True,
        )

    op.drop_table("feishu_users")