"""rename msa_main_agent to policy_agent

Revision ID: a7c1e9b24d6f
Revises: f2c7a1d4e9b0
Create Date: 2026-02-09 12:48:00.000000

"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision: str = 'a7c1e9b24d6f'
down_revision: Union[str, Sequence[str], None] = 'f2c7a1d4e9b0'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def _get_agent_by_name(conn: sa.engine.Connection, name: str):
    return conn.execute(
        sa.text(
            """
            SELECT id, name, display_name
            FROM agent
            WHERE name = :name
            LIMIT 1
            """
        ),
        {"name": name},
    ).mappings().first()


def _get_legacy_policy_agent(conn: sa.engine.Connection):
    return conn.execute(
        sa.text(
            """
            SELECT id, name, display_name
            FROM agent
            WHERE name LIKE 'policy_agent_legacy_%'
            ORDER BY id
            LIMIT 1
            """
        )
    ).mappings().first()


def upgrade() -> None:
    """Upgrade schema."""
    conn = op.get_bind()

    msa_main_agent = _get_agent_by_name(conn, 'msa_main_agent')
    policy_agent = _get_agent_by_name(conn, 'policy_agent')

    # Already migrated or no relevant data to migrate.
    if msa_main_agent is None:
        return

    # If a policy_agent row already exists (e.g., seeded by newer code),
    # move it out of the way so we can preserve existing msa_main_agent row
    # and all profile/version relationships.
    if policy_agent is not None:
        legacy_name = f"policy_agent_legacy_{policy_agent['id']}"
        legacy_display_name = f"Policy Agent Legacy {policy_agent['id']}"
        conn.execute(
            sa.text(
                """
                UPDATE agent
                SET name = :legacy_name,
                    display_name = :legacy_display_name,
                    updated_at = now()
                WHERE id = :agent_id
                """
            ),
            {
                'legacy_name': legacy_name,
                'legacy_display_name': legacy_display_name,
                'agent_id': policy_agent['id'],
            },
        )

    conn.execute(
        sa.text(
            """
            UPDATE agent
            SET name = 'policy_agent',
                display_name = 'Policy Agent',
                updated_at = now()
            WHERE id = :agent_id
            """
        ),
        {'agent_id': msa_main_agent['id']},
    )


def downgrade() -> None:
    """Downgrade schema."""
    conn = op.get_bind()

    msa_main_agent = _get_agent_by_name(conn, 'msa_main_agent')
    policy_agent = _get_agent_by_name(conn, 'policy_agent')

    # Revert canonical rename when possible.
    if policy_agent is not None and msa_main_agent is None:
        conn.execute(
            sa.text(
                """
                UPDATE agent
                SET name = 'msa_main_agent',
                    display_name = 'MSA Main Agent',
                    updated_at = now()
                WHERE id = :agent_id
                """
            ),
            {'agent_id': policy_agent['id']},
        )

    # If upgrade had to move aside an existing policy_agent row,
    # restore it on downgrade.
    legacy = _get_legacy_policy_agent(conn)
    policy_after = _get_agent_by_name(conn, 'policy_agent')
    if legacy is not None and policy_after is None:
        conn.execute(
            sa.text(
                """
                UPDATE agent
                SET name = 'policy_agent',
                    display_name = 'Policy Agent',
                    updated_at = now()
                WHERE id = :agent_id
                """
            ),
            {'agent_id': legacy['id']},
        )
