"""several node keys per account: one per machine the person runs a node on

Revision ID: f8a9b0c1d2e3
Revises: e7f8a9b0c1d2
"""

import uuid
from collections.abc import Sequence
from datetime import UTC, datetime

import sqlalchemy as sa
from alembic import op

revision: str = "f8a9b0c1d2e3"
down_revision: str | Sequence[str] | None = "e7f8a9b0c1d2"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None


def upgrade() -> None:
    node_keys = op.create_table(
        "node_keys",
        sa.Column("id", sa.String(36), primary_key=True),
        sa.Column("user_id", sa.String(36), sa.ForeignKey("users.id"), nullable=False),
        sa.Column("pk_node_ed25519", sa.String(64), nullable=False),
        sa.Column("linked_at", sa.DateTime(timezone=True)),
    )
    op.create_index("uq_node_keys_user_key", "node_keys",
                    ["user_id", "pk_node_ed25519"], unique=True)

    # The key each account had linked keeps working: it becomes that account's
    # first row, so no node has to be linked again after the upgrade.
    conn = op.get_bind()
    linked = conn.execute(sa.text(
        "SELECT id, pk_node_ed25519 FROM users WHERE pk_node_ed25519 IS NOT NULL "
        "AND pk_node_ed25519 <> ''")).fetchall()
    now = datetime.now(UTC)
    if linked:
        op.bulk_insert(node_keys, [
            {"id": str(uuid.uuid4()), "user_id": uid, "pk_node_ed25519": pk,
             "linked_at": now}
            for uid, pk in linked])

    with op.batch_alter_table("users") as batch:
        batch.drop_column("pk_node_ed25519")


def downgrade() -> None:
    with op.batch_alter_table("users") as batch:
        batch.add_column(sa.Column("pk_node_ed25519", sa.String(64), nullable=True))

    # One column holds one key: the most recently linked one is kept, which is
    # the node the account would have had linked last before this revision.
    conn = op.get_bind()
    rows = conn.execute(sa.text(
        "SELECT user_id, pk_node_ed25519 FROM node_keys ORDER BY linked_at")).fetchall()
    latest = {uid: pk for uid, pk in rows}
    for uid, pk in latest.items():
        conn.execute(sa.text("UPDATE users SET pk_node_ed25519 = :pk WHERE id = :uid"),
                     {"pk": pk, "uid": uid})

    op.drop_index("uq_node_keys_user_key", table_name="node_keys")
    op.drop_table("node_keys")