"""add agent zones

Revision ID: 0016_agent_zones
Revises: 0015_miner_control_foundation
Create Date: 2026-08-25
"""

from typing import Sequence, Union

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


revision: str = "0016_agent_zones"
down_revision: Union[str, None] = "0015_miner_control_foundation"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None

DEFAULT_ZONE_NAME = "Default Zone"


def _format_rows(rows: list[tuple]) -> str:
    return ", ".join("/".join(str(value) for value in row) for row in rows)


def upgrade() -> None:
    op.create_table(
        "agent_zones",
        sa.Column(
            "id",
            postgresql.UUID(as_uuid=True),
            primary_key=True,
            server_default=sa.text("gen_random_uuid()"),
        ),
        sa.Column(
            "site_id",
            postgresql.UUID(as_uuid=True),
            sa.ForeignKey("sites.id", ondelete="RESTRICT"),
            nullable=False,
        ),
        sa.Column("name", sa.Text(), nullable=False),
        sa.Column("description", sa.Text(), nullable=True),
        sa.Column(
            "is_enabled", sa.Boolean(), nullable=False, server_default=sa.text("true")
        ),
        sa.Column(
            "created_at",
            sa.DateTime(timezone=True),
            nullable=False,
            server_default=sa.text("now()"),
        ),
        sa.Column(
            "updated_at",
            sa.DateTime(timezone=True),
            nullable=False,
            server_default=sa.text("now()"),
        ),
        sa.CheckConstraint(
            "name = trim(name) AND trim(name) <> '' AND char_length(name) <= 120",
            name="ck_agent_zones_name",
        ),
        sa.UniqueConstraint("id", "site_id", name="uq_agent_zones_id_site_id"),
    )
    op.create_index("idx_agent_zones_site_id", "agent_zones", ["site_id"])
    op.create_index(
        "uq_agent_zones_active_site_name",
        "agent_zones",
        ["site_id", sa.text("lower(name)")],
        unique=True,
        postgresql_where=sa.text("is_enabled"),
    )

    op.create_table(
        "agent_zone_network_ranges",
        sa.Column(
            "id",
            postgresql.UUID(as_uuid=True),
            primary_key=True,
            server_default=sa.text("gen_random_uuid()"),
        ),
        sa.Column(
            "agent_zone_id",
            postgresql.UUID(as_uuid=True),
            sa.ForeignKey("agent_zones.id", ondelete="CASCADE"),
            nullable=False,
        ),
        sa.Column("network", postgresql.CIDR(), nullable=False),
        sa.Column(
            "created_at",
            sa.DateTime(timezone=True),
            nullable=False,
            server_default=sa.text("now()"),
        ),
        sa.UniqueConstraint(
            "agent_zone_id", "network", name="uq_agent_zone_network_ranges_zone_network"
        ),
    )
    op.create_index(
        "idx_agent_zone_network_ranges_zone_id",
        "agent_zone_network_ranges",
        ["agent_zone_id"],
    )

    op.add_column(
        "miners",
        sa.Column("agent_zone_id", postgresql.UUID(as_uuid=True), nullable=True),
    )
    op.create_check_constraint(
        "ck_miners_agent_zone_requires_site",
        "miners",
        "agent_zone_id IS NULL OR site_id IS NOT NULL",
    )
    op.create_foreign_key(
        "fk_miners_agent_zone_site",
        "miners",
        "agent_zones",
        ["agent_zone_id", "site_id"],
        ["id", "site_id"],
        ondelete="RESTRICT",
    )
    op.create_index("idx_miners_agent_zone_id", "miners", ["agent_zone_id"])

    op.add_column(
        "minerops_site_agents",
        sa.Column("agent_zone_id", postgresql.UUID(as_uuid=True), nullable=True),
    )
    op.create_check_constraint(
        "ck_minerops_site_agents_zone_requires_site",
        "minerops_site_agents",
        "agent_zone_id IS NULL OR site_id IS NOT NULL",
    )
    op.create_foreign_key(
        "fk_minerops_site_agents_zone_site",
        "minerops_site_agents",
        "agent_zones",
        ["agent_zone_id", "site_id"],
        ["id", "site_id"],
        ondelete="RESTRICT",
    )

    op.add_column(
        "minerops_agent_snapshots",
        sa.Column("agent_zone_id", postgresql.UUID(as_uuid=True), nullable=True),
    )
    op.create_foreign_key(
        "fk_minerops_agent_snapshots_agent_zone_id",
        "minerops_agent_snapshots",
        "agent_zones",
        ["agent_zone_id"],
        ["id"],
        ondelete="RESTRICT",
    )
    op.create_index(
        "idx_minerops_agent_snapshots_agent_zone_id",
        "minerops_agent_snapshots",
        ["agent_zone_id"],
    )

    op.add_column(
        "minerops_agent_discoveries",
        sa.Column("agent_zone_id", postgresql.UUID(as_uuid=True), nullable=True),
    )
    op.create_foreign_key(
        "fk_minerops_agent_discoveries_agent_zone_id",
        "minerops_agent_discoveries",
        "agent_zones",
        ["agent_zone_id"],
        ["id"],
        ondelete="RESTRICT",
    )
    op.create_index(
        "idx_minerops_agent_discoveries_agent_zone_id",
        "minerops_agent_discoveries",
        ["agent_zone_id"],
    )

    connection = op.get_bind()
    connection.execute(
        sa.text(
            """
            INSERT INTO agent_zones (site_id, name)
            SELECT sites.id, :default_name
            FROM sites
            ORDER BY sites.id
            """
        ),
        {"default_name": DEFAULT_ZONE_NAME},
    )
    connection.execute(
        sa.text(
            """
            UPDATE miners AS miner
            SET agent_zone_id = zone.id
            FROM agent_zones AS zone
            WHERE miner.site_id = zone.site_id
              AND zone.name = :default_name
              AND miner.agent_zone_id IS NULL
            """
        ),
        {"default_name": DEFAULT_ZONE_NAME},
    )
    connection.execute(
        sa.text(
            """
            UPDATE minerops_site_agents AS agent
            SET agent_zone_id = zone.id
            FROM agent_zones AS zone
            WHERE agent.site_id = zone.site_id
              AND zone.name = :default_name
              AND agent.agent_zone_id IS NULL
            """
        ),
        {"default_name": DEFAULT_ZONE_NAME},
    )

    connection.execute(
        sa.text(
            """
            UPDATE minerops_agent_snapshots AS snapshot
            SET agent_zone_id = agent.agent_zone_id
            FROM minerops_site_agents AS agent
            WHERE snapshot.agent_id = agent.agent_id
              AND agent.agent_zone_id IS NOT NULL
              AND snapshot.agent_zone_id IS NULL
            """
        )
    )
    connection.execute(
        sa.text(
            """
            UPDATE minerops_agent_discoveries AS discovery
            SET agent_zone_id = agent.agent_zone_id
            FROM minerops_site_agents AS agent
            WHERE discovery.agent_id = agent.agent_id
              AND agent.agent_zone_id IS NOT NULL
              AND discovery.agent_zone_id IS NULL
            """
        )
    )

    op.create_index(
        "uq_minerops_site_agents_agent_zone_id",
        "minerops_site_agents",
        ["agent_zone_id"],
        unique=True,
        postgresql_where=sa.text("agent_zone_id IS NOT NULL"),
    )
    op.drop_index(
        "uq_minerops_site_agents_site_id", table_name="minerops_site_agents"
    )


def downgrade() -> None:
    connection = op.get_bind()

    duplicate_site_agents = connection.execute(
        sa.text(
            """
            SELECT site_id::text, count(*)
            FROM minerops_site_agents
            WHERE site_id IS NOT NULL
            GROUP BY site_id
            HAVING count(*) > 1
            ORDER BY site_id
            """
        )
    ).all()
    if duplicate_site_agents:
        raise RuntimeError(
            "Cannot downgrade 0016 to 0015: multiple Agent installations are "
            "assigned within Site(s): "
            f"{_format_rows(list(duplicate_site_agents))}. Move or unassign Agents "
            "until each Site has at most one Agent before retrying; no Agent Zone "
            "data was changed."
        )

    non_reducible_zones = connection.execute(
        sa.text(
            """
            SELECT site_id::text, count(*)
            FROM agent_zones
            GROUP BY site_id
            HAVING count(*) > 1
               OR bool_or(name <> :default_name OR description IS NOT NULL OR NOT is_enabled)
            ORDER BY site_id
            """
        ),
        {"default_name": DEFAULT_ZONE_NAME},
    ).all()
    range_count = connection.execute(
        sa.text("SELECT count(*) FROM agent_zone_network_ranges")
    ).scalar_one()
    if non_reducible_zones or range_count:
        detail = _format_rows(list(non_reducible_zones)) or "network ranges exist"
        raise RuntimeError(
            "Cannot downgrade 0016 to 0015: Agent Zone configuration cannot be "
            f"represented by the one-Agent-per-Site model ({detail}). Remove "
            "additional/custom Zones and Zone network ranges before retrying; no "
            "Agent Zone data was changed."
        )

    op.create_index(
        "uq_minerops_site_agents_site_id",
        "minerops_site_agents",
        ["site_id"],
        unique=True,
        postgresql_where=sa.text("site_id IS NOT NULL"),
    )
    op.drop_index(
        "uq_minerops_site_agents_agent_zone_id",
        table_name="minerops_site_agents",
    )

    op.drop_index(
        "idx_minerops_agent_discoveries_agent_zone_id",
        table_name="minerops_agent_discoveries",
    )
    op.drop_constraint(
        "fk_minerops_agent_discoveries_agent_zone_id",
        "minerops_agent_discoveries",
        type_="foreignkey",
    )
    op.drop_column("minerops_agent_discoveries", "agent_zone_id")

    op.drop_index(
        "idx_minerops_agent_snapshots_agent_zone_id",
        table_name="minerops_agent_snapshots",
    )
    op.drop_constraint(
        "fk_minerops_agent_snapshots_agent_zone_id",
        "minerops_agent_snapshots",
        type_="foreignkey",
    )
    op.drop_column("minerops_agent_snapshots", "agent_zone_id")

    op.drop_constraint(
        "fk_minerops_site_agents_zone_site",
        "minerops_site_agents",
        type_="foreignkey",
    )
    op.drop_constraint(
        "ck_minerops_site_agents_zone_requires_site",
        "minerops_site_agents",
        type_="check",
    )
    op.drop_column("minerops_site_agents", "agent_zone_id")

    op.drop_index("idx_miners_agent_zone_id", table_name="miners")
    op.drop_constraint(
        "fk_miners_agent_zone_site", "miners", type_="foreignkey"
    )
    op.drop_constraint(
        "ck_miners_agent_zone_requires_site", "miners", type_="check"
    )
    op.drop_column("miners", "agent_zone_id")

    op.drop_index(
        "idx_agent_zone_network_ranges_zone_id",
        table_name="agent_zone_network_ranges",
    )
    op.drop_table("agent_zone_network_ranges")
    op.drop_index("uq_agent_zones_active_site_name", table_name="agent_zones")
    op.drop_index("idx_agent_zones_site_id", table_name="agent_zones")
    op.drop_table("agent_zones")