"""remove ssh_keys mount from pi-agent manifest

Revision ID: 2026_05_29_remove_ssh_keys_mount_from_manifest
Revises: 2026_05_29_add_ssh_key_ids_to_tool_instances
Create Date: 2026-05-29 14:00:00.000000

"""

import json
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa

# revision identifiers, used by Alembic.
revision: str = "2026_05_29_remove_ssh_keys_mount_from_manifest"
down_revision: Union[str, None] = "2026_05_29_add_ssh_key_ids_to_tool_instances"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    """Remove the ssh_keys mount from the pi-agent manifest."""
    conn = op.get_bind()

    # Get the pi-agent manifest
    result = conn.execute(
        sa.text(
            "SELECT id, manifest FROM tool_definition_manifests WHERE name = 'pi-agent'"
        )
    )
    row = result.fetchone()
    if not row:
        return

    manifest_id, manifest_json = row
    manifest = (
        manifest_json if isinstance(manifest_json, dict) else json.loads(manifest_json)
    )

    mounts = manifest.get("mounts", [])
    original_count = len(mounts)

    # Remove any mount named "ssh_keys"
    filtered_mounts = [m for m in mounts if m.get("name") != "ssh_keys"]

    if len(filtered_mounts) < original_count:
        manifest["mounts"] = filtered_mounts
        conn.execute(
            sa.text(
                "UPDATE tool_definition_manifests SET manifest = :manifest WHERE id = :id"
            ),
            {
                "manifest": json.dumps(manifest),
                "id": manifest_id,
            },
        )


def downgrade() -> None:
    """Restore the ssh_keys mount to the pi-agent manifest."""
    conn = op.get_bind()

    result = conn.execute(
        sa.text(
            "SELECT id, manifest FROM tool_definition_manifests WHERE name = 'pi-agent'"
        )
    )
    row = result.fetchone()
    if not row:
        return

    manifest_id, manifest_json = row
    manifest = (
        manifest_json if isinstance(manifest_json, dict) else json.loads(manifest_json)
    )

    mounts = manifest.get("mounts", [])

    # Check if ssh_keys mount already exists
    if any(m.get("name") == "ssh_keys" for m in mounts):
        return

    # Add the ssh_keys mount back
    mounts.append(
        {
            "name": "ssh_keys",
            "target": "/home/user/.ssh",
            "source_type": "ssh_key",
            "mode": "0700",
            "file_mode": "0600",
            "readonly": True,
        }
    )
    manifest["mounts"] = mounts

    conn.execute(
        sa.text(
            "UPDATE tool_definition_manifests SET manifest = :manifest WHERE id = :id"
        ),
        {
            "manifest": json.dumps(manifest),
            "id": manifest_id,
        },
    )
