"""phase 3 agents

Revision ID: fb27a9b84b1b
Revises: 6b0d8a1c4f2e
"""
import uuid
from collections.abc import Sequence

import sqlalchemy as sa
from alembic import op
from sqlalchemy.dialects import postgresql

revision: str = "fb27a9b84b1b"
down_revision: str | Sequence[str] | None = "6b0d8a1c4f2e"
branch_labels = None
depends_on = None

def upgrade() -> None:
    permissions = sa.table(
        "permissions", sa.column("key", sa.String), sa.column("description", sa.String)
    )
    op.bulk_insert(permissions, [{
        "key": "agents.read", "description": "Read runtime agent configuration"
    }])
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table('agents',
    sa.Column('id', sa.Uuid(), nullable=False),
    sa.Column('workspace_id', sa.Uuid(), nullable=False),
    sa.Column('created_by_user_id', sa.Uuid(), nullable=False),
    sa.Column('name', sa.String(length=160), nullable=False),
    sa.Column('description', sa.String(length=1000), nullable=False),
    sa.Column('purpose', sa.Text(), nullable=False),
    sa.Column('status', sa.String(length=24), nullable=False),
    sa.Column('autonomy_level', sa.Integer(), nullable=False),
    sa.Column('preferred_model', sa.String(length=160), nullable=True),
    sa.Column('fallback_model', sa.String(length=160), nullable=True),
    sa.Column('manager_agent_id', sa.Uuid(), nullable=True),
    sa.Column('escalation_agent_id', sa.Uuid(), nullable=True),
    sa.Column('budget_limit_minor', sa.Integer(), nullable=True),
    sa.Column('budget_currency', sa.String(length=3), nullable=False),
    sa.Column('execution_limit', sa.Integer(), nullable=True),
    sa.Column('version', sa.Integer(), nullable=False),
    sa.Column('created_at', sa.DateTime(timezone=True), server_default=sa.text('(CURRENT_TIMESTAMP)'), nullable=False),
    sa.Column('updated_at', sa.DateTime(timezone=True), server_default=sa.text('(CURRENT_TIMESTAMP)'), nullable=False),
    sa.CheckConstraint("status IN ('active','disabled','archived')", name='ck_agent_status'),
    sa.CheckConstraint('autonomy_level BETWEEN 0 AND 5', name='ck_agent_autonomy'),
    sa.CheckConstraint('budget_limit_minor IS NULL OR budget_limit_minor >= 0', name='ck_agent_budget_limit'),
    sa.CheckConstraint('execution_limit IS NULL OR execution_limit > 0', name='ck_agent_execution_limit'),
    sa.CheckConstraint('version > 0', name='ck_agent_version'),
    sa.ForeignKeyConstraint(['workspace_id', 'created_by_user_id'], ['workspace_memberships.workspace_id', 'workspace_memberships.user_id'], name='fk_agent_creator_membership'),
    sa.ForeignKeyConstraint(['workspace_id', 'escalation_agent_id'], ['agents.workspace_id', 'agents.id'], name='fk_agent_escalation'),
    sa.ForeignKeyConstraint(['workspace_id', 'manager_agent_id'], ['agents.workspace_id', 'agents.id'], name='fk_agent_manager'),
    sa.ForeignKeyConstraint(['workspace_id'], ['workspaces.id'], ),
    sa.PrimaryKeyConstraint('id'),
    sa.UniqueConstraint('workspace_id', 'id', name='uq_agent_workspace_id'),
    sa.UniqueConstraint('workspace_id', 'name', name='uq_agent_workspace_name')
    )
    op.create_index(op.f('ix_agents_created_by_user_id'), 'agents', ['created_by_user_id'], unique=False)
    op.create_index(op.f('ix_agents_workspace_id'), 'agents', ['workspace_id'], unique=False)
    op.create_table('agent_delegations',
    sa.Column('id', sa.Uuid(), nullable=False),
    sa.Column('workspace_id', sa.Uuid(), nullable=False),
    sa.Column('source_agent_id', sa.Uuid(), nullable=False),
    sa.Column('target_agent_id', sa.Uuid(), nullable=False),
    sa.Column('requested_by_user_id', sa.Uuid(), nullable=False),
    sa.Column('instruction', sa.Text(), nullable=False),
    sa.Column('instruction_hash', sa.String(length=64), nullable=False),
    sa.Column('permission_scope', sa.JSON().with_variant(postgresql.JSONB(astext_type=sa.Text()), 'postgresql'), nullable=False),
    sa.Column('status', sa.String(length=24), nullable=False),
    sa.Column('created_at', sa.DateTime(timezone=True), server_default=sa.text('(CURRENT_TIMESTAMP)'), nullable=False),
    sa.Column('completed_at', sa.DateTime(timezone=True), nullable=True),
    sa.CheckConstraint("status IN ('requested','accepted','running','succeeded','failed','cancelled')", name='ck_delegation_status'),
    sa.CheckConstraint('source_agent_id <> target_agent_id', name='ck_delegation_distinct_agents'),
    sa.ForeignKeyConstraint(['workspace_id', 'requested_by_user_id'], ['workspace_memberships.workspace_id', 'workspace_memberships.user_id'], name='fk_delegation_requester'),
    sa.ForeignKeyConstraint(['workspace_id', 'source_agent_id'], ['agents.workspace_id', 'agents.id'], name='fk_delegation_source_agent'),
    sa.ForeignKeyConstraint(['workspace_id', 'target_agent_id'], ['agents.workspace_id', 'agents.id'], name='fk_delegation_target_agent'),
    sa.ForeignKeyConstraint(['workspace_id'], ['workspaces.id'], ),
    sa.PrimaryKeyConstraint('id')
    )
    op.create_index(op.f('ix_agent_delegations_source_agent_id'), 'agent_delegations', ['source_agent_id'], unique=False)
    op.create_index(op.f('ix_agent_delegations_target_agent_id'), 'agent_delegations', ['target_agent_id'], unique=False)
    op.create_index(op.f('ix_agent_delegations_workspace_id'), 'agent_delegations', ['workspace_id'], unique=False)
    op.create_table('agent_permissions',
    sa.Column('workspace_id', sa.Uuid(), nullable=False),
    sa.Column('agent_id', sa.Uuid(), nullable=False),
    sa.Column('permission_key', sa.String(length=120), nullable=False),
    sa.ForeignKeyConstraint(['permission_key'], ['permissions.key'], ),
    sa.ForeignKeyConstraint(['workspace_id', 'agent_id'], ['agents.workspace_id', 'agents.id'], name='fk_agent_permission_agent', ondelete='CASCADE'),
    sa.PrimaryKeyConstraint('workspace_id', 'agent_id', 'permission_key')
    )
    op.create_table('agent_schedules',
    sa.Column('workspace_id', sa.Uuid(), nullable=False),
    sa.Column('agent_id', sa.Uuid(), nullable=False),
    sa.Column('timezone', sa.String(length=80), nullable=False),
    sa.Column('working_days', sa.JSON().with_variant(postgresql.JSONB(astext_type=sa.Text()), 'postgresql'), nullable=False),
    sa.Column('start_minute', sa.Integer(), nullable=False),
    sa.Column('end_minute', sa.Integer(), nullable=False),
    sa.Column('enabled', sa.Boolean(), nullable=False),
    sa.CheckConstraint('end_minute BETWEEN 1 AND 1440', name='ck_agent_schedule_end'),
    sa.CheckConstraint('start_minute < end_minute', name='ck_agent_schedule_range'),
    sa.CheckConstraint('start_minute BETWEEN 0 AND 1439', name='ck_agent_schedule_start'),
    sa.ForeignKeyConstraint(['workspace_id', 'agent_id'], ['agents.workspace_id', 'agents.id'], name='fk_agent_schedule_agent', ondelete='CASCADE'),
    sa.PrimaryKeyConstraint('workspace_id', 'agent_id')
    )
    op.create_table('agent_tools',
    sa.Column('workspace_id', sa.Uuid(), nullable=False),
    sa.Column('agent_id', sa.Uuid(), nullable=False),
    sa.Column('tool_name', sa.String(length=120), nullable=False),
    sa.Column('constraints', sa.JSON().with_variant(postgresql.JSONB(astext_type=sa.Text()), 'postgresql'), nullable=False),
    sa.ForeignKeyConstraint(['workspace_id', 'agent_id'], ['agents.workspace_id', 'agents.id'], name='fk_agent_tool_agent', ondelete='CASCADE'),
    sa.PrimaryKeyConstraint('workspace_id', 'agent_id', 'tool_name')
    )
    op.create_table('agent_versions',
    sa.Column('id', sa.Uuid(), nullable=False),
    sa.Column('workspace_id', sa.Uuid(), nullable=False),
    sa.Column('agent_id', sa.Uuid(), nullable=False),
    sa.Column('version_number', sa.Integer(), nullable=False),
    sa.Column('system_instructions', sa.Text(), nullable=False),
    sa.Column('personality', sa.Text(), nullable=False),
    sa.Column('config_snapshot', sa.JSON().with_variant(postgresql.JSONB(astext_type=sa.Text()), 'postgresql'), nullable=False),
    sa.Column('created_by_user_id', sa.Uuid(), nullable=False),
    sa.Column('created_at', sa.DateTime(timezone=True), server_default=sa.text('(CURRENT_TIMESTAMP)'), nullable=False),
    sa.CheckConstraint('version_number > 0', name='ck_agent_version_number'),
    sa.ForeignKeyConstraint(['workspace_id', 'agent_id'], ['agents.workspace_id', 'agents.id'], name='fk_agent_version_agent', ondelete='CASCADE'),
    sa.ForeignKeyConstraint(['workspace_id', 'created_by_user_id'], ['workspace_memberships.workspace_id', 'workspace_memberships.user_id'], name='fk_agent_version_creator'),
    sa.PrimaryKeyConstraint('id'),
    sa.UniqueConstraint('workspace_id', 'agent_id', 'version_number', name='uq_agent_version_number')
    )
    op.create_index(op.f('ix_agent_versions_agent_id'), 'agent_versions', ['agent_id'], unique=False)
    op.create_index(op.f('ix_agent_versions_workspace_id'), 'agent_versions', ['workspace_id'], unique=False)
    _seed_agent_read_permission()
    _backfill_default_agents()
    if op.get_bind().dialect.name == "sqlite":
        with op.batch_alter_table("executions", recreate="always") as batch:
            batch.add_column(sa.Column("agent_id", sa.Uuid(), nullable=True))
            batch.create_index("ix_executions_agent_id", ["agent_id"])
            batch.create_foreign_key(
                "fk_execution_agent",
                "agents",
                ["workspace_id", "agent_id"],
                ["workspace_id", "id"],
            )
    else:
        op.add_column("executions", sa.Column("agent_id", sa.Uuid(), nullable=True))
        op.create_index("ix_executions_agent_id", "executions", ["agent_id"])
        op.create_foreign_key(
            "fk_execution_agent",
            "executions",
            "agents",
            ["workspace_id", "agent_id"],
            ["workspace_id", "id"],
        )
    # ### end Alembic commands ###


def _seed_agent_read_permission() -> None:
    roles = sa.table(
        "roles",
        sa.column("id", sa.Uuid),
        sa.column("workspace_id", sa.Uuid),
        sa.column("name", sa.String),
        sa.column("is_system", sa.Boolean),
    )
    role_permissions = sa.table(
        "role_permissions",
        sa.column("workspace_id", sa.Uuid),
        sa.column("role_id", sa.Uuid),
        sa.column("permission_key", sa.String),
    )
    op.get_bind().execute(
        role_permissions.insert().from_select(
            ["workspace_id", "role_id", "permission_key"],
            sa.select(roles.c.workspace_id, roles.c.id, sa.literal("agents.read")).where(
                roles.c.is_system.is_(True),
                roles.c.name.in_([
                    "Owner", "Administrator", "Assistant Manager", "Human Assistant",
                    "Viewer", "Developer",
                ]),
            ),
        )
    )


def _backfill_default_agents() -> None:
    roles = sa.table(
        "roles", sa.column("id", sa.Uuid), sa.column("workspace_id", sa.Uuid),
        sa.column("name", sa.String), sa.column("is_system", sa.Boolean),
    )
    membership_roles = sa.table(
        "membership_roles", sa.column("workspace_id", sa.Uuid),
        sa.column("membership_id", sa.Uuid), sa.column("role_id", sa.Uuid),
    )
    memberships = sa.table(
        "workspace_memberships", sa.column("id", sa.Uuid),
        sa.column("workspace_id", sa.Uuid), sa.column("user_id", sa.Uuid),
        sa.column("status", sa.String),
    )
    agents = sa.table(
        "agents", sa.column("id", sa.Uuid), sa.column("workspace_id", sa.Uuid),
        sa.column("created_by_user_id", sa.Uuid), sa.column("name", sa.String),
        sa.column("description", sa.String), sa.column("purpose", sa.Text),
        sa.column("status", sa.String), sa.column("autonomy_level", sa.Integer),
        sa.column("preferred_model", sa.String), sa.column("fallback_model", sa.String),
        sa.column("manager_agent_id", sa.Uuid), sa.column("escalation_agent_id", sa.Uuid),
        sa.column("budget_limit_minor", sa.Integer), sa.column("budget_currency", sa.String),
        sa.column("execution_limit", sa.Integer), sa.column("version", sa.Integer),
    )
    versions = sa.table(
        "agent_versions", sa.column("id", sa.Uuid), sa.column("workspace_id", sa.Uuid),
        sa.column("agent_id", sa.Uuid), sa.column("version_number", sa.Integer),
        sa.column("system_instructions", sa.Text), sa.column("personality", sa.Text),
        sa.column("config_snapshot", sa.JSON), sa.column("created_by_user_id", sa.Uuid),
    )
    agent_permissions = sa.table(
        "agent_permissions", sa.column("workspace_id", sa.Uuid),
        sa.column("agent_id", sa.Uuid), sa.column("permission_key", sa.String),
    )
    agent_tools = sa.table(
        "agent_tools", sa.column("workspace_id", sa.Uuid),
        sa.column("agent_id", sa.Uuid), sa.column("tool_name", sa.String),
        sa.column("constraints", sa.JSON),
    )
    connection = op.get_bind()
    owners = connection.execute(
        sa.select(roles.c.workspace_id, memberships.c.user_id)
        .select_from(
            roles.join(
                membership_roles,
                (membership_roles.c.workspace_id == roles.c.workspace_id)
                & (membership_roles.c.role_id == roles.c.id),
            ).join(
                memberships,
                (memberships.c.workspace_id == membership_roles.c.workspace_id)
                & (memberships.c.id == membership_roles.c.membership_id),
            )
        )
        .where(
            roles.c.name == "Owner", roles.c.is_system.is_(True),
            memberships.c.status == "active",
        )
    ).all()
    default_permissions = [
        "calendar.read", "email.read", "whatsapp.read", "tasks.read", "memory.read",
        "workspace.profile.read", "system.status.read",
    ]
    default_tools = ["workspace.profile.read", "system.status.read"]
    for workspace_id, user_id in owners:
        agent_id = uuid.uuid4()
        snapshot = {
            "name": "Personal Executive Assistant", "status": "active",
            "autonomy_level": 3, "tool_names": default_tools,
            "permission_keys": default_permissions, "schedule": None,
        }
        connection.execute(agents.insert().values(
            id=agent_id, workspace_id=workspace_id, created_by_user_id=user_id,
            name="Personal Executive Assistant",
            description="Primary workspace assistant",
            purpose="Act as the user's primary AI assistant within explicit permissions.",
            status="active", autonomy_level=3, preferred_model=None, fallback_model=None,
            manager_agent_id=None, escalation_agent_id=None, budget_limit_minor=None,
            budget_currency="USD", execution_limit=None, version=1,
        ))
        connection.execute(versions.insert().values(
            id=uuid.uuid4(), workspace_id=workspace_id, agent_id=agent_id,
            version_number=1,
            system_instructions="Assist the user without exceeding granted tools or permissions.",
            personality="Helpful, concise, careful, and transparent about uncertainty.",
            config_snapshot=snapshot, created_by_user_id=user_id,
        ))
        connection.execute(agent_permissions.insert(), [
            {"workspace_id": workspace_id, "agent_id": agent_id, "permission_key": key}
            for key in default_permissions
        ])
        connection.execute(agent_tools.insert(), [
            {"workspace_id": workspace_id, "agent_id": agent_id, "tool_name": name,
             "constraints": {}}
            for name in default_tools
        ])

def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    if op.get_bind().dialect.name == "sqlite":
        with op.batch_alter_table("executions", recreate="always") as batch:
            batch.drop_constraint("fk_execution_agent", type_="foreignkey")
            batch.drop_index("ix_executions_agent_id")
            batch.drop_column("agent_id")
    else:
        op.drop_constraint("fk_execution_agent", "executions", type_="foreignkey")
        op.drop_index("ix_executions_agent_id", table_name="executions")
        op.drop_column("executions", "agent_id")
    op.drop_index(op.f('ix_agent_versions_workspace_id'), table_name='agent_versions')
    op.drop_index(op.f('ix_agent_versions_agent_id'), table_name='agent_versions')
    op.drop_table('agent_versions')
    op.drop_table('agent_tools')
    op.drop_table('agent_schedules')
    op.drop_table('agent_permissions')
    op.drop_index(op.f('ix_agent_delegations_workspace_id'), table_name='agent_delegations')
    op.drop_index(op.f('ix_agent_delegations_target_agent_id'), table_name='agent_delegations')
    op.drop_index(op.f('ix_agent_delegations_source_agent_id'), table_name='agent_delegations')
    op.drop_table('agent_delegations')
    op.drop_index(op.f('ix_agents_workspace_id'), table_name='agents')
    op.drop_index(op.f('ix_agents_created_by_user_id'), table_name='agents')
    op.drop_table('agents')
    role_permissions = sa.table(
        "role_permissions", sa.column("permission_key", sa.String)
    )
    permissions = sa.table("permissions", sa.column("key", sa.String))
    connection = op.get_bind()
    connection.execute(role_permissions.delete().where(
        role_permissions.c.permission_key == "agents.read"
    ))
    connection.execute(permissions.delete().where(permissions.c.key == "agents.read"))
    # ### end Alembic commands ###