"""add experiment project_id column

Revision ID: add_experiment_project_id
Revises: add_budget_tables
Create Date: 2026-02-25
"""

from alembic import op
import sqlalchemy as sa

revision = "add_experiment_project_id"
down_revision = "add_budget_tables"
branch_labels = None
depends_on = None


def _has_column(table_name: str, column_name: str) -> bool:
    bind = op.get_bind()
    inspector = sa.inspect(bind)
    columns = inspector.get_columns(table_name)
    return any(col["name"] == column_name for col in columns)


def _has_index(table_name: str, index_name: str) -> bool:
    bind = op.get_bind()
    inspector = sa.inspect(bind)
    indexes = inspector.get_indexes(table_name)
    return any(idx["name"] == index_name for idx in indexes)


def upgrade():
    if not _has_column("experiments", "project_id"):
        op.add_column("experiments", sa.Column("project_id", sa.BigInteger(), nullable=True))

    bind = op.get_bind()
    dialect = bind.dialect.name

    # 尽量从 linked_context_nodes 推断 project_id（仅唯一归属时回填）
    if dialect == "sqlite":
        op.execute(
            """
            UPDATE experiments
            SET project_id = (
                SELECT tn.project_id
                FROM truth_nodes tn
                WHERE instr(COALESCE(experiments.linked_context_nodes, ''), printf('%s', tn.id)) > 0
                  AND tn.project_id IS NOT NULL
                GROUP BY tn.project_id
                HAVING COUNT(*) >= 1
                LIMIT 1
            )
            WHERE project_id IS NULL
            """
        )
    else:
        # PostgreSQL JSON/数组形态不稳定，迁移阶段不做激进回填，留给服务层按需补写。
        pass

    if not _has_index("experiments", "ix_experiments_project_id"):
        op.create_index("ix_experiments_project_id", "experiments", ["project_id"])


def downgrade():
    if _has_index("experiments", "ix_experiments_project_id"):
        op.drop_index("ix_experiments_project_id", table_name="experiments")
    if _has_column("experiments", "project_id"):
        op.drop_column("experiments", "project_id")