"""add kb_imports and kb_import_nodes tables

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

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


def upgrade():
    op.create_table(
        "kb_imports",
        sa.Column("id", UUID(as_uuid=True), primary_key=True, server_default=sa.text("gen_random_uuid()")),
        sa.Column("account_id", UUID(as_uuid=True), sa.ForeignKey("accounts.id", ondelete="CASCADE"), nullable=False, index=True),
        sa.Column("created_by", UUID(as_uuid=True), sa.ForeignKey("users.id", ondelete="CASCADE"), nullable=False, index=True),
        sa.Column("source_filename", sa.String(500), nullable=True),
        sa.Column("source_format", sa.String(20), nullable=False),
        sa.Column("source_text", sa.Text, nullable=False),
        sa.Column("source_metadata", JSONB, nullable=True),
        sa.Column("target_type", sa.String(20), nullable=False),
        sa.Column("status", sa.String(20), nullable=False, server_default="processing"),
        sa.Column("confidence_avg", sa.Float, nullable=True),
        sa.Column("error_message", sa.Text, nullable=True),
        sa.Column("processing_time_ms", sa.Integer, nullable=True),
        sa.Column("ai_tokens_input", sa.Integer, nullable=True),
        sa.Column("ai_tokens_output", sa.Integer, nullable=True),
        sa.Column("tree_id", UUID(as_uuid=True), sa.ForeignKey("trees.id", ondelete="SET NULL"), nullable=True),
        sa.Column("batch_id", UUID(as_uuid=True), nullable=True, index=True),
        sa.Column("created_at", sa.DateTime(timezone=True), nullable=False, server_default=sa.text("now()")),
        sa.Column("updated_at", sa.DateTime(timezone=True), nullable=False, server_default=sa.text("now()")),
        sa.CheckConstraint(
            "source_format IN ('txt', 'paste', 'docx', 'pdf', 'html', 'md')",
            name="ck_kb_imports_source_format",
        ),
        sa.CheckConstraint(
            "target_type IN ('troubleshooting', 'procedural')",
            name="ck_kb_imports_target_type",
        ),
        sa.CheckConstraint(
            "status IN ('processing', 'ready', 'committed', 'failed')",
            name="ck_kb_imports_status",
        ),
    )

    op.create_index("ix_kb_imports_status", "kb_imports", ["status"])
    op.create_index("ix_kb_imports_created_at_desc", "kb_imports", [sa.text("created_at DESC")])

    op.create_table(
        "kb_import_nodes",
        sa.Column("id", UUID(as_uuid=True), primary_key=True, server_default=sa.text("gen_random_uuid()")),
        sa.Column("kb_import_id", UUID(as_uuid=True), sa.ForeignKey("kb_imports.id", ondelete="CASCADE"), nullable=False, index=True),
        sa.Column("node_order", sa.Integer, nullable=False),
        sa.Column("node_type", sa.String(20), nullable=False),
        sa.Column("content", JSONB, nullable=False),
        sa.Column("parent_node_id", UUID(as_uuid=True), sa.ForeignKey("kb_import_nodes.id", ondelete="SET NULL"), nullable=True),
        sa.Column("source_excerpt", sa.Text, nullable=True),
        sa.Column("confidence_score", sa.Float, nullable=False),
        sa.Column("user_edited", sa.Boolean, nullable=False, server_default=sa.text("false")),
        sa.Column("user_approved", sa.Boolean, nullable=False, server_default=sa.text("false")),
        sa.Column("created_at", sa.DateTime(timezone=True), nullable=False, server_default=sa.text("now()")),
        sa.Column("updated_at", sa.DateTime(timezone=True), nullable=False, server_default=sa.text("now()")),
        sa.CheckConstraint(
            "node_type IN ('question', 'resolution', 'step', 'section_header', 'warning', 'action')",
            name="ck_kb_import_nodes_node_type",
        ),
    )

    op.create_index("ix_kb_import_nodes_confidence", "kb_import_nodes", ["confidence_score"])


def downgrade():
    op.drop_table("kb_import_nodes")
    op.drop_table("kb_imports")