"""add agentkit extension tables

Revision ID: b001_agentkit_extension
Revises: a79329c23b20
Create Date: 2026-06-04

新增 fischer-agentkit 扩展表：
- episodic_memories: 经验记忆
- evolution_logs: 进化日志
- ab_test_configs: A/B测试配置
"""

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

revision = "b001_agentkit_extension"
down_revision = "a79329c23b20"
branch_labels = None
depends_on = None


def upgrade() -> None:
    # episodic_memories
    op.create_table(
        "episodic_memories",
        sa.Column("id", sa.Uuid(), nullable=False),
        sa.Column("agent_name", sa.String(50), nullable=False),
        sa.Column("task_type", sa.String(50), nullable=False),
        sa.Column("input_data", postgresql.JSONB(), nullable=True),
        sa.Column("output_data", postgresql.JSONB(), nullable=True),
        sa.Column("success", sa.Boolean(), nullable=True),
        sa.Column("quality_score", sa.Float(), nullable=True),
        sa.Column("reflection", sa.Text(), nullable=True),
        sa.Column("embedding_id", sa.String(100), nullable=True),
        sa.Column("tags", postgresql.JSONB(), nullable=True),
        sa.Column("created_at", sa.DateTime(timezone=True), server_default=sa.func.now(), nullable=False),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_index("idx_episodic_agent_name", "episodic_memories", ["agent_name"])
    op.create_index("idx_episodic_task_type", "episodic_memories", ["task_type"])
    op.create_index("idx_episodic_created_at", "episodic_memories", ["created_at"])

    # evolution_logs
    op.create_table(
        "evolution_logs",
        sa.Column("id", sa.Uuid(), nullable=False),
        sa.Column("agent_name", sa.String(50), nullable=False),
        sa.Column("change_type", sa.String(30), nullable=False),
        sa.Column("before", postgresql.JSONB(), nullable=True),
        sa.Column("after", postgresql.JSONB(), nullable=True),
        sa.Column("metrics", postgresql.JSONB(), nullable=True),
        sa.Column("ab_test_id", sa.String(100), nullable=True),
        sa.Column("status", sa.String(20), server_default="pending", nullable=False),
        sa.Column("created_at", sa.DateTime(timezone=True), server_default=sa.func.now(), nullable=False),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_index("idx_evolution_agent_name", "evolution_logs", ["agent_name"])
    op.create_index("idx_evolution_change_type", "evolution_logs", ["change_type"])
    op.create_index("idx_evolution_status", "evolution_logs", ["status"])
    op.create_index("idx_evolution_created_at", "evolution_logs", ["created_at"])

    # ab_test_configs
    op.create_table(
        "ab_test_configs",
        sa.Column("id", sa.Uuid(), nullable=False),
        sa.Column("agent_name", sa.String(50), nullable=False),
        sa.Column("test_name", sa.String(100), nullable=False),
        sa.Column("variant_a", postgresql.JSONB(), nullable=False),
        sa.Column("variant_b", postgresql.JSONB(), nullable=False),
        sa.Column("traffic_split", sa.Float(), server_default="0.5", nullable=False),
        sa.Column("status", sa.String(20), server_default="running", nullable=False),
        sa.Column("winner", sa.String(10), nullable=True),
        sa.Column("metrics", postgresql.JSONB(), nullable=True),
        sa.Column("started_at", sa.DateTime(timezone=True), server_default=sa.func.now(), nullable=False),
        sa.Column("ended_at", sa.DateTime(timezone=True), nullable=True),
        sa.Column("created_at", sa.DateTime(timezone=True), server_default=sa.func.now(), nullable=False),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_index("idx_ab_test_agent_name", "ab_test_configs", ["agent_name"])
    op.create_index("idx_ab_test_status", "ab_test_configs", ["status"])


def downgrade() -> None:
    op.drop_table("ab_test_configs")
    op.drop_table("evolution_logs")
    op.drop_table("episodic_memories")