"""Hooks for revision 88b0fb10906f (webhooks and queue updates)."""

from __future__ import annotations

import json

from nanoid import generate as generate_nanoid
from sqlalchemy import text

from tests.alembic.registry import register_after_upgrade, register_before_upgrade
from tests.alembic.verifier import MigrationVerifier

PAYLOAD_MARKER = generate_nanoid()
WORKSPACE_NAME = generate_nanoid()
SESSION_NAME = generate_nanoid()
SENDER_NAME = generate_nanoid()
TARGET_NAME = generate_nanoid()


@register_before_upgrade("88b0fb10906f")
def prepare_webhooks(verifier: MigrationVerifier) -> None:
    # Validate objects introduced by the migration are absent
    verifier.assert_table_exists("webhook_endpoints", exists=False)
    verifier.assert_column_exists("queue", "task_type", exists=False)
    verifier.assert_column_exists("queue", "work_unit_key", exists=False)

    verifier.assert_column_exists(
        "active_queue_sessions", "work_unit_key", exists=False
    )
    verifier.assert_constraint_exists(
        "active_queue_sessions", "unique_work_unit_key", "unique", exists=False
    )

    # Validate objects removed by the migration still exist
    verifier.assert_column_exists("active_queue_sessions", "session_id")
    verifier.assert_column_exists("active_queue_sessions", "sender_name")
    verifier.assert_column_exists("active_queue_sessions", "target_name")
    verifier.assert_column_exists("active_queue_sessions", "task_type")
    verifier.assert_constraint_exists(
        "active_queue_sessions", "unique_active_queue_session", "unique"
    )

    # Bulk insert 500k queue items to test batched backfill
    verifier.conn.execute(text("SET LOCAL synchronous_commit = OFF"))
    verifier.conn.execute(
        text(
            f'INSERT INTO "{verifier.schema}"."queue" ("payload") '
            + "SELECT jsonb_build_object("
            + "  'marker', 'bulk-' || substr(md5(random()::text), 1, 21),"
            + "  'task_type', 'representation',"
            + f"  'workspace_name', '{WORKSPACE_NAME}',"
            + f"  'session_name', '{SESSION_NAME}',"
            + f"  'sender_name', '{SENDER_NAME}',"
            + f"  'target_name', '{TARGET_NAME}'"
            + ") FROM generate_series(1, :n)"
        ),
        {"n": 500_000},
    )

    payload = json.dumps(
        {
            "marker": PAYLOAD_MARKER,
            "task_type": "representation",
            "workspace_name": WORKSPACE_NAME,
            "session_name": SESSION_NAME,
            "sender_name": SENDER_NAME,
            "target_name": TARGET_NAME,
        }
    )

    # Add single queue item to test expected data transformation
    verifier.conn.execute(
        text(f'INSERT INTO "{verifier.schema}"."queue" ("payload") VALUES (:payload)'),
        {"payload": payload},
    )


@register_after_upgrade("88b0fb10906f")
def verify_webhooks_and_queue(verifier: MigrationVerifier) -> None:
    """Validate queue enrichment and webhook table creation."""

    verifier.assert_table_exists("webhook_endpoints")

    verifier.assert_column_exists("queue", "task_type", nullable=False)
    verifier.assert_column_exists("queue", "work_unit_key", nullable=False)
    verifier.assert_column_exists("active_queue_sessions", "work_unit_key")
    verifier.assert_constraint_exists(
        "active_queue_sessions", "unique_work_unit_key", "unique"
    )

    verifier.assert_column_exists("active_queue_sessions", "session_id", exists=False)
    verifier.assert_column_exists("active_queue_sessions", "sender_name", exists=False)
    verifier.assert_column_exists("active_queue_sessions", "target_name", exists=False)
    verifier.assert_column_exists("active_queue_sessions", "task_type", exists=False)
    verifier.assert_constraint_exists(
        "active_queue_sessions", "unique_active_queue_session", "unique", exists=False
    )

    # Ensure no NULLs remain after backfill across all rows
    verifier.assert_no_nulls("queue", "task_type")
    verifier.assert_no_nulls("queue", "work_unit_key")

    row = verifier.conn.execute(
        text(
            f'SELECT "task_type", "work_unit_key" FROM "{verifier.schema}"."queue" '
            + "WHERE payload->>'marker' = :marker"
        ),
        {"marker": PAYLOAD_MARKER},
    ).one()

    assert row.task_type == "representation"
    expected_key = (
        f"representation:{WORKSPACE_NAME}:{SESSION_NAME}:{SENDER_NAME}:{TARGET_NAME}"
    )
    assert row.work_unit_key == expected_key