"""add KB Accelerator columns to plan_limits

Revision ID: 055
Revises: 054
"""
from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects.postgresql import JSONB

revision = "055"
down_revision = "054"
branch_labels = None
depends_on = None


def upgrade():
    # Add KB Accelerator columns to plan_limits
    op.add_column("plan_limits", sa.Column("kb_accelerator_enabled", sa.Boolean, nullable=False, server_default=sa.text("false")))
    op.add_column("plan_limits", sa.Column("kb_max_lifetime_conversions", sa.Integer, nullable=True))
    op.add_column("plan_limits", sa.Column("kb_batch_max_size", sa.Integer, nullable=True))
    op.add_column("plan_limits", sa.Column("kb_allowed_formats", JSONB, nullable=False, server_default=sa.text("'[\"txt\",\"paste\"]'::jsonb")))
    op.add_column("plan_limits", sa.Column("kb_detailed_analysis", sa.Boolean, nullable=False, server_default=sa.text("false")))
    op.add_column("plan_limits", sa.Column("kb_conversational_refinement", sa.Boolean, nullable=False, server_default=sa.text("false")))
    op.add_column("plan_limits", sa.Column("kb_step_library_matching", sa.Boolean, nullable=False, server_default=sa.text("false")))
    op.add_column("plan_limits", sa.Column("kb_history_limit", sa.Integer, nullable=True))

    # Seed defaults for each plan tier
    op.execute("""
        UPDATE plan_limits SET
            kb_accelerator_enabled = true,
            kb_max_lifetime_conversions = 3,
            kb_batch_max_size = NULL,
            kb_allowed_formats = '["txt","paste"]'::jsonb,
            kb_detailed_analysis = false,
            kb_conversational_refinement = false,
            kb_step_library_matching = false,
            kb_history_limit = 3
        WHERE plan = 'free'
    """)

    op.execute("""
        UPDATE plan_limits SET
            kb_accelerator_enabled = true,
            kb_max_lifetime_conversions = NULL,
            kb_batch_max_size = 5,
            kb_allowed_formats = '["txt","paste","docx","pdf","html","md"]'::jsonb,
            kb_detailed_analysis = true,
            kb_conversational_refinement = true,
            kb_step_library_matching = true,
            kb_history_limit = NULL
        WHERE plan = 'pro'
    """)

    op.execute("""
        UPDATE plan_limits SET
            kb_accelerator_enabled = true,
            kb_max_lifetime_conversions = NULL,
            kb_batch_max_size = 10,
            kb_allowed_formats = '["txt","paste","docx","pdf","html","md"]'::jsonb,
            kb_detailed_analysis = true,
            kb_conversational_refinement = true,
            kb_step_library_matching = true,
            kb_history_limit = NULL
        WHERE plan = 'team'
    """)


def downgrade():
    op.drop_column("plan_limits", "kb_history_limit")
    op.drop_column("plan_limits", "kb_step_library_matching")
    op.drop_column("plan_limits", "kb_conversational_refinement")
    op.drop_column("plan_limits", "kb_detailed_analysis")
    op.drop_column("plan_limits", "kb_allowed_formats")
    op.drop_column("plan_limits", "kb_batch_max_size")
    op.drop_column("plan_limits", "kb_max_lifetime_conversions")
    op.drop_column("plan_limits", "kb_accelerator_enabled")