"""adopt tenant-safe historical case learning tables

Revision ID: 20260716_0008
Revises: 20260716_0007
Create Date: 2026-07-16 13:20:00
"""

from __future__ import annotations

from collections.abc import Sequence

import sqlalchemy as sa

from alembic import op

revision: str = "20260716_0008"
down_revision: str | None = "20260716_0007"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None


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


def _inspector() -> sa.Inspector:
    return sa.inspect(op.get_bind())


def _table_exists(table_name: str) -> bool:
    return table_name in _inspector().get_table_names()


def _columns(table_name: str) -> set[str]:
    return {str(item["name"]) for item in _inspector().get_columns(table_name)}


def _indexes(table_name: str) -> dict[str, dict[str, object]]:
    return {
        str(item["name"]): item for item in _inspector().get_indexes(table_name) if item.get("name")
    }


def _unique_constraints(table_name: str) -> dict[str, dict[str, object]]:
    return {
        str(item["name"]): item
        for item in _inspector().get_unique_constraints(table_name)
        if item.get("name")
    }


def _foreign_keys(table_name: str) -> list[dict[str, object]]:
    return list(_inspector().get_foreign_keys(table_name))


def _ensure_index(
    table_name: str,
    index_name: str,
    columns: list[str],
    *,
    unique: bool = False,
) -> None:
    if index_name not in _indexes(table_name):
        op.create_index(index_name, table_name, columns, unique=unique)


def _drop_legacy_single_key_uniqueness(table_name: str, key_column: str) -> None:
    for name, item in list(_unique_constraints(table_name).items()):
        if tuple(item.get("column_names") or ()) == (key_column,):
            op.drop_constraint(name, table_name, type_="unique")
    # PostgreSQL 会把唯一约束的支撑索引同时返回给 get_indexes；先删除约束并
    # 重新反射，避免对支撑索引执行 DROP INDEX 导致 dependency error。
    for name, item in list(_indexes(table_name).items()):
        if tuple(item.get("column_names") or ()) == (key_column,) and bool(item.get("unique")):
            op.drop_index(name, table_name=table_name)


def _ensure_composite_unique(
    table_name: str,
    constraint_name: str,
    columns: list[str],
) -> None:
    expected = tuple(columns)
    if any(
        tuple(item.get("column_names") or ()) == expected
        for item in _unique_constraints(table_name).values()
    ):
        return
    op.create_unique_constraint(constraint_name, table_name, columns)


def _drop_index_if_exists(table_name: str, index_name: str) -> None:
    if index_name in _indexes(table_name):
        op.drop_index(index_name, table_name=table_name)


def _drop_unique_if_exists(table_name: str, constraint_name: str) -> None:
    if constraint_name in _unique_constraints(table_name):
        op.drop_constraint(constraint_name, table_name, type_="unique")


def _require_lossless_default_tenant_downgrade(table_name: str) -> None:
    non_default_count = int(
        op.get_bind().scalar(
            sa.text(
                f"SELECT COUNT(*) FROM {table_name} "
                "WHERE tenant_id IS DISTINCT FROM 'default'"
            )
        )
        or 0
    )
    if non_default_count:
        raise RuntimeError(
            f"cannot downgrade {table_name}: non-default tenant data would lose isolation"
        )


def _require_lossless_few_shot_downgrade() -> None:
    enriched_count = int(
        op.get_bind().scalar(
            sa.text(
                "SELECT COUNT(*) FROM few_shot_samples "
                "WHERE COALESCE(TRIM(policy_ref), '') <> '' "
                "OR COALESCE(TRIM(rule_version), '') <> ''"
            )
        )
        or 0
    )
    if enriched_count:
        raise RuntimeError(
            "cannot downgrade few_shot_samples: "
            f"{enriched_count} row(s) contain policy_ref or rule_version data"
        )


def _drop_risk_observation_claim_foreign_keys() -> None:
    """统一采用软引用，避免 Head 结构取决于 expense_claims 的创建时机。"""
    for item in _foreign_keys("risk_observations"):
        if (
            tuple(item.get("constrained_columns") or ()) == ("claim_id",)
            and item.get("referred_table") == "expense_claims"
            and item.get("name")
        ):
            op.drop_constraint(
                str(item["name"]),
                "risk_observations",
                type_="foreignkey",
            )


def _create_risk_observations() -> None:
    op.create_table(
        "risk_observations",
        sa.Column("id", sa.String(length=36), nullable=False),
        sa.Column("tenant_id", sa.String(length=64), server_default="default", nullable=False),
        sa.Column("observation_key", sa.String(length=160), nullable=False),
        sa.Column("subject_type", sa.String(length=50), nullable=False),
        sa.Column("subject_key", sa.String(length=160), nullable=False),
        sa.Column("subject_label", sa.String(length=160), nullable=False, server_default=""),
        # expense_claims 仍由 legacy bootstrap 创建，空库迁移到此版本时并不存在。
        # claim_id 因此是显式软引用，与 ORM 的 viewonly relationship 保持一致。
        sa.Column("claim_id", sa.String(length=36), nullable=True),
        sa.Column("claim_no", sa.String(length=80), nullable=False, server_default=""),
        sa.Column("run_id", sa.String(length=80), nullable=True),
        sa.Column("execution_log_id", sa.String(length=36), nullable=True),
        sa.Column("risk_type", sa.String(length=80), nullable=False),
        sa.Column("risk_signal", sa.String(length=100), nullable=False),
        sa.Column("title", sa.String(length=200), nullable=False, server_default=""),
        sa.Column("description", sa.Text(), nullable=False, server_default=""),
        sa.Column("risk_score", sa.Integer(), nullable=False, server_default="0"),
        sa.Column("risk_level", sa.String(length=20), nullable=False),
        sa.Column("confidence_score", sa.Float(), nullable=False, server_default="0"),
        sa.Column("control_stage", sa.String(length=50), nullable=False, server_default=""),
        sa.Column("control_mode", sa.String(length=50), nullable=False, server_default=""),
        sa.Column("automation_mode", sa.String(length=50), nullable=False, server_default=""),
        sa.Column("source", sa.String(length=60), nullable=False, server_default=""),
        sa.Column("algorithm_version", sa.String(length=80), nullable=False, server_default=""),
        sa.Column("status", sa.String(length=30), nullable=False, server_default="pending_review"),
        sa.Column(
            "feedback_status", sa.String(length=30), nullable=False, server_default="unreviewed"
        ),
        sa.Column("contribution_scores_json", sa.JSON(), nullable=False),
        sa.Column("baseline_json", sa.JSON(), nullable=False),
        sa.Column("evidence_json", sa.JSON(), nullable=False),
        sa.Column("graph_node_keys_json", sa.JSON(), nullable=False),
        sa.Column("graph_edge_keys_json", sa.JSON(), nullable=False),
        sa.Column("policy_refs_json", sa.JSON(), nullable=False),
        sa.Column("similar_case_claim_ids_json", sa.JSON(), nullable=False),
        sa.Column("ontology_json", sa.JSON(), nullable=False),
        sa.Column("decision_trace_json", sa.JSON(), nullable=False),
        sa.Column(
            "created_at", sa.DateTime(timezone=True), server_default=sa.func.now(), nullable=False
        ),
        sa.Column(
            "updated_at", sa.DateTime(timezone=True), server_default=sa.func.now(), nullable=False
        ),
        sa.PrimaryKeyConstraint("id"),
        sa.UniqueConstraint("tenant_id", "observation_key", name="uq_risk_observations_tenant_key"),
    )


def _adopt_risk_observations() -> None:
    _drop_risk_observation_claim_foreign_keys()
    if "tenant_id" not in _columns("risk_observations"):
        op.add_column(
            "risk_observations",
            sa.Column("tenant_id", sa.String(length=64), server_default="default", nullable=False),
        )
    op.execute(
        sa.text(
            "UPDATE risk_observations SET tenant_id = 'default' "
            "WHERE tenant_id IS NULL OR tenant_id = ''"
        )
    )
    _drop_legacy_single_key_uniqueness("risk_observations", "observation_key")
    _ensure_composite_unique(
        "risk_observations",
        "uq_risk_observations_tenant_key",
        ["tenant_id", "observation_key"],
    )


def _create_risk_observation_feedback() -> None:
    op.create_table(
        "risk_observation_feedback",
        sa.Column("id", sa.String(length=36), nullable=False),
        sa.Column("observation_id", sa.String(length=36), nullable=False),
        sa.Column("feedback_type", sa.String(length=30), nullable=False),
        sa.Column("action", sa.String(length=50), nullable=False, server_default=""),
        sa.Column("actor", sa.String(length=100), nullable=False, server_default=""),
        sa.Column("comment", sa.Text(), nullable=True),
        sa.Column("payload_json", sa.JSON(), nullable=False),
        sa.Column(
            "created_at", sa.DateTime(timezone=True), server_default=sa.func.now(), nullable=False
        ),
        sa.ForeignKeyConstraint(["observation_id"], ["risk_observations.id"]),
        sa.PrimaryKeyConstraint("id"),
    )


def _create_few_shot_samples() -> None:
    op.create_table(
        "few_shot_samples",
        sa.Column("id", sa.String(length=36), nullable=False),
        sa.Column("tenant_id", sa.String(length=64), server_default="default", nullable=False),
        sa.Column("sample_key", sa.String(length=160), nullable=False),
        sa.Column("source_observation_id", sa.String(length=36), nullable=True),
        sa.Column(
            "scene", sa.String(length=50), nullable=False, server_default="risk_rule_generation"
        ),
        sa.Column("policy_ref", sa.String(length=160), nullable=False, server_default=""),
        sa.Column("rule_version", sa.String(length=80), nullable=False, server_default=""),
        sa.Column("domain", sa.String(length=50), nullable=False, server_default=""),
        sa.Column("risk_type", sa.String(length=80), nullable=False, server_default=""),
        sa.Column("risk_level", sa.String(length=20), nullable=False, server_default=""),
        sa.Column("label", sa.String(length=30), nullable=False, server_default="confirmed"),
        sa.Column("case_text", sa.Text(), nullable=False, server_default=""),
        sa.Column("conclusion_text", sa.Text(), nullable=False, server_default=""),
        sa.Column("payload_json", sa.JSON(), nullable=False),
        sa.Column("vector_id", sa.String(length=100), nullable=True),
        sa.Column("status", sa.String(length=20), nullable=False, server_default="active"),
        sa.Column("created_at", sa.DateTime(), server_default=sa.func.now(), nullable=False),
        sa.Column("updated_at", sa.DateTime(), server_default=sa.func.now(), nullable=False),
        sa.ForeignKeyConstraint(["source_observation_id"], ["risk_observations.id"]),
        sa.PrimaryKeyConstraint("id"),
        sa.UniqueConstraint("tenant_id", "sample_key", name="uq_few_shot_samples_tenant_key"),
    )


def _adopt_few_shot_samples() -> None:
    existing = _columns("few_shot_samples")
    for column in (
        sa.Column("tenant_id", sa.String(length=64), server_default="default", nullable=False),
        sa.Column("policy_ref", sa.String(length=160), server_default="", nullable=False),
        sa.Column("rule_version", sa.String(length=80), server_default="", nullable=False),
    ):
        if column.name not in existing:
            op.add_column("few_shot_samples", column)
    op.execute(
        sa.text(
            "UPDATE few_shot_samples SET tenant_id = 'default' "
            "WHERE tenant_id IS NULL OR tenant_id = ''"
        )
    )
    op.execute(sa.text("UPDATE few_shot_samples SET policy_ref = '' WHERE policy_ref IS NULL"))
    op.execute(sa.text("UPDATE few_shot_samples SET rule_version = '' WHERE rule_version IS NULL"))
    _drop_legacy_single_key_uniqueness("few_shot_samples", "sample_key")
    _ensure_composite_unique(
        "few_shot_samples",
        "uq_few_shot_samples_tenant_key",
        ["tenant_id", "sample_key"],
    )


def _ensure_risk_indexes() -> None:
    definitions = {
        "ix_risk_observations_tenant_id": ["tenant_id"],
        "ix_risk_observations_observation_key": ["observation_key"],
        "ix_risk_observations_subject_type": ["subject_type"],
        "ix_risk_observations_subject_key": ["subject_key"],
        "ix_risk_observations_claim_id": ["claim_id"],
        "ix_risk_observations_claim_no": ["claim_no"],
        "ix_risk_observations_run_id": ["run_id"],
        "ix_risk_observations_execution_log_id": ["execution_log_id"],
        "ix_risk_observations_risk_type": ["risk_type"],
        "ix_risk_observations_risk_signal": ["risk_signal"],
        "ix_risk_observations_risk_score": ["risk_score"],
        "ix_risk_observations_risk_level": ["risk_level"],
        "ix_risk_observations_source": ["source"],
        "ix_risk_observations_algorithm_version": ["algorithm_version"],
        "ix_risk_observations_status": ["status"],
        "ix_risk_observations_feedback_status": ["feedback_status"],
        "ix_risk_observations_subject": ["subject_type", "subject_key"],
        "ix_risk_observations_signal_level": ["risk_signal", "risk_level"],
        "ix_risk_observations_status_created": ["status", "created_at"],
        "ix_risk_observations_tenant_status": ["tenant_id", "status", "created_at"],
    }
    for name, columns in definitions.items():
        _ensure_index("risk_observations", name, columns)
    _ensure_index(
        "risk_observation_feedback",
        "ix_risk_observation_feedback_observation_id",
        ["observation_id"],
    )
    _ensure_index(
        "risk_observation_feedback", "ix_risk_observation_feedback_feedback_type", ["feedback_type"]
    )
    _ensure_index(
        "risk_observation_feedback",
        "ix_risk_observation_feedback_type_created",
        ["feedback_type", "created_at"],
    )


def _ensure_few_shot_indexes() -> None:
    definitions = {
        "ix_few_shot_samples_tenant_id": ["tenant_id"],
        "ix_few_shot_samples_sample_key": ["sample_key"],
        "ix_few_shot_samples_source_observation_id": ["source_observation_id"],
        "ix_few_shot_samples_scene": ["scene"],
        "ix_few_shot_samples_policy_ref": ["policy_ref"],
        "ix_few_shot_samples_rule_version": ["rule_version"],
        "ix_few_shot_samples_domain": ["domain"],
        "ix_few_shot_samples_risk_type": ["risk_type"],
        "ix_few_shot_samples_label": ["label"],
        "ix_few_shot_samples_status": ["status"],
        "ix_few_shot_samples_scene_label": ["scene", "label"],
        "ix_few_shot_samples_domain_risk_type": ["domain", "risk_type"],
        "ix_few_shot_samples_tenant_rule_lookup": [
            "tenant_id",
            "scene",
            "policy_ref",
            "rule_version",
            "status",
        ],
    }
    for name, columns in definitions.items():
        _ensure_index("few_shot_samples", name, columns)


def upgrade() -> None:
    _require_postgresql()
    if _table_exists("risk_observations"):
        _adopt_risk_observations()
    else:
        _create_risk_observations()

    if not _table_exists("risk_observation_feedback"):
        _create_risk_observation_feedback()

    if _table_exists("few_shot_samples"):
        _adopt_few_shot_samples()
    else:
        _create_few_shot_samples()

    _ensure_risk_indexes()
    _ensure_few_shot_indexes()


def downgrade() -> None:
    _require_postgresql()
    # 0008 之前这三张表由旧 bootstrap 管理。降级必须恢复旧结构而不是删除
    # 已确认的风险反馈与 few-shot 数据；若存在非默认租户数据则无法无损合并，
    # 直接失败并保持当前版本，禁止用数据丢失换取“成功降级”。
    few_shot_exists = _table_exists("few_shot_samples")
    risk_observations_exists = _table_exists("risk_observations")
    # 所有数据安全检查必须先于 DDL；任一检查失败时完整保留 0008 结构。
    if few_shot_exists:
        _require_lossless_default_tenant_downgrade("few_shot_samples")
        _require_lossless_few_shot_downgrade()
    if risk_observations_exists:
        _require_lossless_default_tenant_downgrade("risk_observations")

    if few_shot_exists:
        for index_name in (
            "ix_few_shot_samples_tenant_rule_lookup",
            "ix_few_shot_samples_tenant_id",
            "ix_few_shot_samples_policy_ref",
            "ix_few_shot_samples_rule_version",
        ):
            _drop_index_if_exists("few_shot_samples", index_name)
        _drop_unique_if_exists(
            "few_shot_samples",
            "uq_few_shot_samples_tenant_key",
        )
        op.create_unique_constraint(
            "uq_few_shot_samples_sample_key",
            "few_shot_samples",
            ["sample_key"],
        )
        op.drop_column("few_shot_samples", "rule_version")
        op.drop_column("few_shot_samples", "policy_ref")
        op.drop_column("few_shot_samples", "tenant_id")

    if risk_observations_exists:
        for index_name in (
            "ix_risk_observations_tenant_status",
            "ix_risk_observations_tenant_id",
        ):
            _drop_index_if_exists("risk_observations", index_name)
        _drop_unique_if_exists(
            "risk_observations",
            "uq_risk_observations_tenant_key",
        )
        op.create_unique_constraint(
            "uq_risk_observations_observation_key",
            "risk_observations",
            ["observation_key"],
        )
        op.drop_column("risk_observations", "tenant_id")
