"""fix_pi_agent_home_directory_mount

Revision ID: 2026_06_14_182955
Revises: fc8f1a20cbf6
Create Date: 2026-06-14 18:29:55.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_14_182955"
down_revision: Union[str, Sequence[str], None] = "fc8f1a20cbf6"
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

    # Mount the repo under the configured home directory, preserving the repo
    # directory name via the WORKSPACE_NAME runtime variable.
    for mount in manifest.get("mounts", []):
        if mount.get("source_type") == "repo":
            mount["target"] = "~/{{WORKSPACE_NAME}}"

    # Keep /workspace as a compatibility symlink to the real mount path.
    runtime = manifest.setdefault("runtime", {})
    runtime["working_dir"] = "/workspace"

    # Update the startup script to chown the real mount path.
    scripts = manifest.setdefault("scripts", {})
    scripts["startup"] = [
        'if [ -n "$WORKSPACE_NAME" ]; then sudo chown -R user:user "$HOME/$WORKSPACE_NAME" 2>/dev/null || true; fi',
    ]

    _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

    for mount in manifest.get("mounts", []):
        if mount.get("source_type") == "repo":
            mount["target"] = "/workspace"

    runtime = manifest.setdefault("runtime", {})
    runtime["working_dir"] = "/workspace"

    scripts = manifest.setdefault("scripts", {})
    scripts["startup"] = [
        "if [ -d /workspace ]; then sudo chown -R user:user /workspace 2>/dev/null || true; fi",
    ]

    _update_manifest(conn, manifest_id, manifest)
