"""remove pi agent explicit repo mount

Revision ID: 2026_06_15_090500
Revises: 2026_06_14_182955
Create Date: 2026-06-15 09:05:00.000000

"""

from typing import Sequence, Union

import sqlalchemy as sa
from alembic import op
from sqlalchemy.sql import column, table

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

tool_definition_manifests = table(
    "tool_definition_manifests",
    column("id", sa.UUID),
    column("name", sa.String),
    column("manifest", sa.JSON),
)


def _find_pi_agent_manifest(
    conn: sa.Connection,
) -> tuple[Union[str, None], Union[dict, None]]:
    result = conn.execute(
        sa.select(
            tool_definition_manifests.c.id, tool_definition_manifests.c.manifest
        ).where(tool_definition_manifests.c.name == "pi-agent")
    ).fetchone()
    if result is None:
        return None, None
    return result.id, dict(result.manifest)


def _update_manifest(conn: sa.Connection, manifest_id: str, manifest: dict) -> None:
    conn.execute(
        sa.update(tool_definition_manifests)
        .where(tool_definition_manifests.c.id == manifest_id)
        .values(manifest=manifest)
    )


def upgrade() -> None:
    conn = op.get_bind()
    manifest_id, manifest = _find_pi_agent_manifest(conn)
    if not manifest_id or not manifest:
        return

    # The repo mount is now synthesized by compile_compose based on the
    # instance's repository, so the manifest no longer needs an explicit
    # repo mount with a {{WORKSPACE_NAME}} placeholder.
    manifest["mounts"] = [
        mount
        for mount in manifest.get("mounts", [])
        if mount.get("source_type") != "repo"
    ]

    _update_manifest(conn, manifest_id, manifest)


def downgrade() -> None:
    conn = op.get_bind()
    manifest_id, manifest = _find_pi_agent_manifest(conn)
    if not manifest_id or not manifest:
        return

    mounts = manifest.setdefault("mounts", [])
    if not any(mount.get("source_type") == "repo" for mount in mounts):
        mounts.append(
            {
                "name": "workspace",
                "target": "~/{{WORKSPACE_NAME}}",
                "source_type": "repo",
                "writable": True,
                "owner": "user",
            }
        )

    _update_manifest(conn, manifest_id, manifest)
