"""add structural tenant and platform scope to Agent assets

Revision ID: 20260717_0026
Revises: 20260717_0025
Create Date: 2026-07-17 14:30:00
"""

from collections.abc import Sequence

import sqlalchemy as sa

from alembic import op

revision: str = "20260717_0026"
down_revision: str | None = "20260717_0025"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None

_ASSET_TABLE = "agent_assets"
_ASSET_CHILD_TABLES = (
    "agent_asset_versions",
    "agent_asset_reviews",
    "agent_asset_test_runs",
    "agent_asset_rule_feedback",
)
_OWNERSHIP_CHILD_TABLES = (
    "agent_asset_versions",
    "agent_asset_reviews",
)
_ONLYOFFICE_SESSION_TABLE = "agent_asset_onlyoffice_sessions"


def _require_postgresql() -> None:
    dialect_name = op.get_bind().dialect.name
    if dialect_name != "postgresql":
        raise RuntimeError(
            "20260717_0026 only supports PostgreSQL; "
            f"refusing to mutate {dialect_name} without transactional constraint DDL"
        )


def _has_complete_agent_asset_schema() -> bool:
    inspector = sa.inspect(op.get_bind())
    table_names = (_ASSET_TABLE, *_ASSET_CHILD_TABLES)
    existing = {table_name for table_name in table_names if inspector.has_table(table_name)}
    if not existing:
        return False
    missing = set(table_names) - existing
    if missing:
        raise RuntimeError(
            "cannot migrate partial Agent asset schema; missing tables: "
            + ", ".join(sorted(missing))
        )
    return True


def _create_onlyoffice_session_table() -> None:
    if sa.inspect(op.get_bind()).has_table(_ONLYOFFICE_SESSION_TABLE):
        return
    op.create_table(
        _ONLYOFFICE_SESSION_TABLE,
        sa.Column("jti", sa.String(length=36), nullable=False),
        sa.Column("tenant_id", sa.String(length=64), nullable=False),
        sa.Column("resource_scope", sa.String(length=16), nullable=False),
        sa.Column("asset_id", sa.String(length=100), nullable=False),
        sa.Column("document_key", sa.String(length=200), nullable=False),
        sa.Column("document_version", sa.String(length=30), nullable=False),
        sa.Column("document_fingerprint", sa.String(length=160), nullable=False),
        sa.Column("audience", sa.String(length=80), nullable=False),
        sa.Column("writable", sa.Boolean(), nullable=False),
        sa.Column("status", sa.String(length=16), nullable=False),
        sa.Column("actor", sa.String(length=160), nullable=False),
        sa.Column("expires_at", sa.DateTime(timezone=True), nullable=False),
        sa.Column("claimed_at", sa.DateTime(timezone=True), nullable=True),
        sa.Column("consumed_at", sa.DateTime(timezone=True), nullable=True),
        sa.Column("failure_reason", sa.Text(), nullable=False),
        sa.Column(
            "created_at",
            sa.DateTime(timezone=True),
            nullable=False,
            server_default=sa.func.now(),
        ),
        sa.CheckConstraint(
            "(resource_scope = 'platform' AND tenant_id = 'platform') OR "
            "(resource_scope = 'tenant' AND tenant_id <> 'platform')",
            name="ck_agent_asset_onlyoffice_sessions_scope_tenant",
        ),
        sa.CheckConstraint(
            "status IN ('active', 'processing', 'consumed', 'failed', 'revoked')",
            name="ck_agent_asset_onlyoffice_sessions_status",
        ),
        sa.CheckConstraint(
            "(status = 'active' AND claimed_at IS NULL AND consumed_at IS NULL) OR "
            "(status IN ('processing', 'failed') AND claimed_at IS NOT NULL "
            "AND consumed_at IS NULL) OR "
            "(status = 'consumed' AND claimed_at IS NOT NULL AND consumed_at IS NOT NULL) OR "
            "(status = 'revoked' AND consumed_at IS NULL)",
            name="ck_agent_asset_onlyoffice_sessions_lifecycle",
        ),
        sa.ForeignKeyConstraint(
            ["tenant_id"],
            ["tenants.tenant_id"],
            name="fk_agent_asset_onlyoffice_sessions_tenant",
            ondelete="CASCADE",
        ),
        sa.PrimaryKeyConstraint("jti"),
    )
    op.create_index(
        "ix_agent_asset_onlyoffice_sessions_tenant_asset",
        _ONLYOFFICE_SESSION_TABLE,
        ["tenant_id", "resource_scope", "asset_id", "created_at"],
    )
    op.create_index(
        "ix_agent_asset_onlyoffice_sessions_status_expiry",
        _ONLYOFFICE_SESSION_TABLE,
        ["status", "expires_at"],
    )


def _add_scope_columns(table_name: str) -> None:
    op.add_column(table_name, sa.Column("tenant_id", sa.String(length=64), nullable=True))
    op.add_column(table_name, sa.Column("scope", sa.String(length=16), nullable=True))


def _assert_known_asset_tenants() -> None:
    unknown = int(
        op.get_bind().scalar(
            sa.text(
                """
                SELECT COUNT(*)
                FROM agent_assets AS asset
                LEFT JOIN tenants AS tenant
                  ON tenant.tenant_id = NULLIF(BTRIM(asset.config_json ->> 'tenant_id'), '')
                WHERE NULLIF(BTRIM(asset.config_json ->> 'tenant_id'), '') IS NOT NULL
                  AND tenant.tenant_id IS NULL
                """
            )
        )
        or 0
    )
    if unknown:
        raise RuntimeError(
            "cannot migrate Agent assets: config_json contains tenant ids absent from tenants "
            f"(unknown_assets={unknown})"
        )


def _backfill_scope() -> None:
    _assert_known_asset_tenants()
    op.execute(
        """
        UPDATE agent_assets
        SET tenant_id = COALESCE(
                NULLIF(BTRIM(config_json ->> 'tenant_id'), ''),
                'platform'
            ),
            scope = CASE
                WHEN NULLIF(BTRIM(config_json ->> 'tenant_id'), '') IS NULL
                    THEN 'platform'
                ELSE 'tenant'
            END
        """
    )
    for table_name in _OWNERSHIP_CHILD_TABLES:
        op.execute(
            f"""
            UPDATE {table_name} AS child
            SET tenant_id = asset.tenant_id,
                scope = asset.scope
            FROM agent_assets AS asset
            WHERE child.asset_id = asset.id
            """
        )
    op.execute(
        """
        UPDATE agent_asset_test_runs AS child
        SET tenant_id = COALESCE(
                NULLIF(BTRIM(child.input_json ->> 'target_tenant_id'), ''),
                NULLIF(BTRIM(asset.config_json ->> 'tenant_id'), ''),
                asset.tenant_id
            ),
            scope = CASE
                WHEN COALESCE(
                    NULLIF(BTRIM(child.input_json ->> 'target_tenant_id'), ''),
                    NULLIF(BTRIM(asset.config_json ->> 'tenant_id'), ''),
                    asset.tenant_id
                ) = 'platform' THEN 'platform'
                ELSE 'tenant'
            END
        FROM agent_assets AS asset
        WHERE child.asset_id = asset.id
        """
    )
    op.execute(
        """
        UPDATE agent_asset_rule_feedback AS child
        SET tenant_id = COALESCE(
                NULLIF(BTRIM(asset.config_json ->> 'tenant_id'), ''),
                asset.tenant_id
            ),
            scope = CASE
                WHEN COALESCE(
                    NULLIF(BTRIM(asset.config_json ->> 'tenant_id'), ''),
                    asset.tenant_id
                ) = 'platform' THEN 'platform'
                ELSE 'tenant'
            END
        FROM agent_assets AS asset
        WHERE child.asset_id = asset.id
        """
    )


def _set_scope_not_null_and_defaults(table_name: str) -> None:
    op.alter_column(
        table_name,
        "tenant_id",
        existing_type=sa.String(length=64),
        nullable=False,
        server_default="platform",
    )
    op.alter_column(
        table_name,
        "scope",
        existing_type=sa.String(length=16),
        nullable=False,
        server_default="platform",
    )


def _assert_known_scoped_tenants() -> None:
    bind = op.get_bind()
    for table_name in (_ASSET_TABLE, *_ASSET_CHILD_TABLES):
        unknown = int(
            bind.scalar(
                sa.text(
                    f"""
                    SELECT COUNT(*)
                    FROM {table_name} AS scoped
                    LEFT JOIN tenants AS tenant ON tenant.tenant_id = scoped.tenant_id
                    WHERE tenant.tenant_id IS NULL
                    """
                )
            )
            or 0
        )
        if unknown:
            raise RuntimeError(
                f"cannot migrate {table_name}: scoped tenant is absent from tenants "
                f"(unknown_rows={unknown})"
            )


def _create_scope_check(table_name: str) -> None:
    op.create_check_constraint(
        f"ck_{table_name}_scope_tenant",
        table_name,
        "(scope = 'platform' AND tenant_id = 'platform') OR "
        "(scope = 'tenant' AND tenant_id <> 'platform')",
    )


def _drop_legacy_code_uniqueness() -> None:
    op.execute("ALTER TABLE agent_assets DROP CONSTRAINT IF EXISTS agent_assets_code_key")
    op.execute("DROP INDEX IF EXISTS ix_agent_assets_code")
    op.create_index("ix_agent_assets_code", _ASSET_TABLE, ["code"], unique=False)


def upgrade() -> None:
    _require_postgresql()
    has_agent_asset_schema = _has_complete_agent_asset_schema()
    _create_onlyoffice_session_table()
    if not has_agent_asset_schema:
        return
    _add_scope_columns(_ASSET_TABLE)
    for table_name in _ASSET_CHILD_TABLES:
        _add_scope_columns(table_name)
    _backfill_scope()
    _set_scope_not_null_and_defaults(_ASSET_TABLE)
    for table_name in _ASSET_CHILD_TABLES:
        _set_scope_not_null_and_defaults(table_name)
    _assert_known_scoped_tenants()

    _drop_legacy_code_uniqueness()
    op.create_unique_constraint(
        "uq_agent_assets_tenant_scope_id",
        _ASSET_TABLE,
        ["tenant_id", "scope", "id"],
    )
    op.create_unique_constraint(
        "uq_agent_assets_tenant_scope_code",
        _ASSET_TABLE,
        ["tenant_id", "scope", "code"],
    )
    op.create_foreign_key(
        "fk_agent_assets_tenant",
        _ASSET_TABLE,
        "tenants",
        ["tenant_id"],
        ["tenant_id"],
        ondelete="RESTRICT",
    )
    _create_scope_check(_ASSET_TABLE)
    op.create_index(
        "ix_agent_assets_scope_tenant",
        _ASSET_TABLE,
        ["scope", "tenant_id"],
    )

    for table_name in _ASSET_CHILD_TABLES:
        op.create_foreign_key(
            f"fk_{table_name}_tenant",
            table_name,
            "tenants",
            ["tenant_id"],
            ["tenant_id"],
            ondelete="RESTRICT",
        )
        _create_scope_check(table_name)
        op.create_index(
            f"ix_{table_name}_tenant_asset",
            table_name,
            ["tenant_id", "scope", "asset_id"],
        )
    for table_name in _OWNERSHIP_CHILD_TABLES:
        op.create_foreign_key(
            f"fk_{table_name}_tenant_asset",
            table_name,
            _ASSET_TABLE,
            ["tenant_id", "scope", "asset_id"],
            ["tenant_id", "scope", "id"],
            ondelete="CASCADE",
        )


def _require_lossless_downgrade() -> None:
    bind = op.get_bind()
    tenant_assets = int(
        bind.scalar(
            sa.text("SELECT COUNT(*) FROM agent_assets WHERE scope = 'tenant'")
        )
        or 0
    )
    tenant_evidence = sum(
        int(
            bind.scalar(
                sa.text(f"SELECT COUNT(*) FROM {table_name} WHERE scope = 'tenant'")
            )
            or 0
        )
        for table_name in ("agent_asset_test_runs", "agent_asset_rule_feedback")
    )
    if tenant_assets or tenant_evidence:
        raise RuntimeError(
            "cannot downgrade Agent asset tenant security: tenant-owned facts exist "
            f"(tenant_assets={tenant_assets}, tenant_evidence={tenant_evidence})"
        )


def _require_no_onlyoffice_sessions() -> None:
    if not sa.inspect(op.get_bind()).has_table(_ONLYOFFICE_SESSION_TABLE):
        return
    session_count = int(
        op.get_bind().scalar(
            sa.text(f"SELECT COUNT(*) FROM {_ONLYOFFICE_SESSION_TABLE}")
        )
        or 0
    )
    if session_count:
        raise RuntimeError(
            "cannot downgrade Agent asset tenant security: ONLYOFFICE session evidence "
            f"exists (sessions={session_count})"
        )


def _drop_onlyoffice_session_table() -> None:
    if not sa.inspect(op.get_bind()).has_table(_ONLYOFFICE_SESSION_TABLE):
        return
    op.drop_index(
        "ix_agent_asset_onlyoffice_sessions_status_expiry",
        table_name=_ONLYOFFICE_SESSION_TABLE,
    )
    op.drop_index(
        "ix_agent_asset_onlyoffice_sessions_tenant_asset",
        table_name=_ONLYOFFICE_SESSION_TABLE,
    )
    op.drop_table(_ONLYOFFICE_SESSION_TABLE)


def _drop_constraint_if_exists(
    table_name: str,
    constraint_name: str,
) -> None:
    """删除本迁移负责的约束，兼容模型建表产生的不同外键名称。"""
    preparer = op.get_bind().dialect.identifier_preparer
    op.execute(
        sa.text(
            f"ALTER TABLE {preparer.quote(table_name)} "
            f"DROP CONSTRAINT IF EXISTS {preparer.quote(constraint_name)}"
        )
    )


def downgrade() -> None:
    _require_postgresql()
    has_agent_asset_schema = _has_complete_agent_asset_schema()
    _require_no_onlyoffice_sessions()
    _drop_onlyoffice_session_table()
    if not has_agent_asset_schema:
        return
    _require_lossless_downgrade()
    for table_name in reversed(_OWNERSHIP_CHILD_TABLES):
        _drop_constraint_if_exists(
            table_name,
            f"fk_{table_name}_tenant_asset",
        )
    for table_name in reversed(_ASSET_CHILD_TABLES):
        _drop_constraint_if_exists(
            table_name,
            f"fk_{table_name}_tenant",
        )
        op.drop_index(f"ix_{table_name}_tenant_asset", table_name=table_name)
        _drop_constraint_if_exists(
            table_name,
            f"ck_{table_name}_scope_tenant",
        )
    op.drop_index("ix_agent_assets_scope_tenant", table_name=_ASSET_TABLE)
    _drop_constraint_if_exists(_ASSET_TABLE, "ck_agent_assets_scope_tenant")
    _drop_constraint_if_exists(_ASSET_TABLE, "fk_agent_assets_tenant")
    _drop_constraint_if_exists(
        _ASSET_TABLE,
        "uq_agent_assets_tenant_scope_code",
    )
    _drop_constraint_if_exists(
        _ASSET_TABLE,
        "uq_agent_assets_tenant_scope_id",
    )
    op.drop_index("ix_agent_assets_code", table_name=_ASSET_TABLE)
    op.create_index("ix_agent_assets_code", _ASSET_TABLE, ["code"], unique=True)
    for table_name in reversed(_ASSET_CHILD_TABLES):
        op.drop_column(table_name, "scope")
        op.drop_column(table_name, "tenant_id")
    op.drop_column(_ASSET_TABLE, "scope")
    op.drop_column(_ASSET_TABLE, "tenant_id")