"""remove dkuCallerTicket from existing merged_context rows

Revision ID: 7b9f8c1d2e3a
Revises: 2368ae4f97cc
Create Date: 2026-04-02 12:00:00.000000
"""

import json
import os
import sys
from typing import Sequence, Union

import sqlalchemy as sa

# Add alembic folder to path for utils import
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from utils import get_qualified_table_name, timed_step

# revision identifiers, used by Alembic.
revision: str = "7b9f8c1d2e3a"
down_revision: Union[str, Sequence[str], None] = "2368ae4f97cc"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def remove_dku_caller_ticket(value):
    if value is None:
        return None

    try:
        merged_context = json.loads(value)
    except Exception:
        return value

    if not isinstance(merged_context, dict) or "dkuCallerTicket" not in merged_context:
        return value

    merged_context.pop("dkuCallerTicket", None)
    return json.dumps(merged_context)


def upgrade() -> None:
    """Remove dkuCallerTicket from previously stored message merged contexts."""
    from alembic import op

    conn = op.get_bind()
    table_name = get_qualified_table_name("messages")
    placeholder = "%s" if conn.dialect.name in ("postgresql", "snowflake") else "?"

    with timed_step("Backfill messages.merged_context"):
        rows = conn.execute(
            sa.text(f"SELECT id, merged_context FROM {table_name} WHERE merged_context IS NOT NULL")
        )
        for message_id, merged_context_raw in rows:
            updated_merged_context = remove_dku_caller_ticket(merged_context_raw)
            if updated_merged_context == merged_context_raw:
                continue

            conn.exec_driver_sql(
                f"UPDATE {table_name} SET merged_context = {placeholder} WHERE id = {placeholder}",
                (updated_merged_context, message_id),
            )


def downgrade() -> None:
    """Downgrade schema."""
    pass
