"""update profile model and add rag models

Revision ID: 3c72cb3be254
Revises: 766207a9c1b1
Create Date: 2025-03-25 08:56:11.571652

"""

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

# revision identifiers, used by Alembic.
revision = "3c72cb3be254"
down_revision = "b7b8d3df0452"
branch_labels = None
depends_on = None


def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table(
        "knowledge_base",
        sa.Column("id", mysql.BIGINT(), autoincrement=True, nullable=False),
        sa.Column(
            "kb_id", sa.String(length=36), nullable=False, comment="Knowledge Base ID"
        ),
        sa.Column(
            "kb_name",
            sa.String(length=255),
            nullable=False,
            comment="Knowledge Base name",
        ),
        sa.Column(
            "kb_description",
            sa.Text(),
            nullable=True,
            comment="Knowledge Base description",
        ),
        sa.Column(
            "embedding_model",
            sa.String(length=255),
            nullable=False,
            comment="Embedding Model name",
        ),
        sa.Column("dim", sa.Integer(), nullable=False, comment="VectorDB dim"),
        sa.Column(
            "tag_ids",
            sa.Text(),
            nullable=True,
            comment="Knowledge Base Tag IDs, separated by ','",
        ),
        sa.Column(
            "course_ids",
            sa.Text(),
            nullable=True,
            comment="Course UUIDs, separated by ','",
        ),
        sa.Column("meta_data", sa.Text(), nullable=True, comment="Meta Data"),
        sa.Column("extra_data", sa.Text(), nullable=True, comment="Extra Data"),
        sa.Column(
            "created_user_id",
            sa.String(length=36),
            nullable=True,
            comment="created user ID",
        ),
        sa.Column(
            "updated_user_id",
            sa.String(length=36),
            nullable=True,
            comment="updated user ID",
        ),
        sa.Column("created", sa.TIMESTAMP(), nullable=False, comment="Creation time"),
        sa.Column("updated", sa.TIMESTAMP(), nullable=False, comment="Update time"),
        sa.PrimaryKeyConstraint("id"),
    )
    with op.batch_alter_table("knowledge_base", schema=None) as batch_op:
        batch_op.create_index(
            batch_op.f("ix_knowledge_base_kb_id"), ["kb_id"], unique=False
        )

    op.create_table(
        "knowledge_chunk",
        sa.Column("id", mysql.BIGINT(), autoincrement=True, nullable=False),
        sa.Column(
            "kb_id", sa.String(length=36), nullable=False, comment="Knowledge Base ID"
        ),
        sa.Column("file_id", sa.String(length=36), nullable=False, comment="File ID"),
        sa.Column("chunk_id", sa.String(length=64), nullable=False, comment="Chunk ID"),
        sa.Column("chunk_index", sa.Integer(), nullable=True, comment="Chunk index"),
        sa.Column("chunk_text", sa.Text(), nullable=True, comment="Chunk text"),
        sa.Column("chunk_vector", sa.Text(), nullable=True, comment="Chunk vector"),
        sa.Column("meta_data", sa.Text(), nullable=True, comment="Meta Data"),
        sa.Column("extra_data", sa.Text(), nullable=True, comment="Extra Data"),
        sa.Column(
            "created_user_id",
            sa.String(length=36),
            nullable=True,
            comment="created user ID",
        ),
        sa.Column(
            "updated_user_id",
            sa.String(length=36),
            nullable=True,
            comment="updated user ID",
        ),
        sa.Column("created", sa.TIMESTAMP(), nullable=False, comment="Creation time"),
        sa.Column("updated", sa.TIMESTAMP(), nullable=False, comment="Update time"),
        sa.PrimaryKeyConstraint("id"),
    )
    with op.batch_alter_table("knowledge_chunk", schema=None) as batch_op:
        batch_op.create_index(
            batch_op.f("ix_knowledge_chunk_chunk_id"), ["chunk_id"], unique=False
        )
        batch_op.create_index(
            batch_op.f("ix_knowledge_chunk_file_id"), ["file_id"], unique=False
        )
        batch_op.create_index(
            batch_op.f("ix_knowledge_chunk_kb_id"), ["kb_id"], unique=False
        )

    op.create_table(
        "knowledge_file",
        sa.Column("id", mysql.BIGINT(), autoincrement=True, nullable=False),
        sa.Column(
            "kb_id", sa.String(length=36), nullable=False, comment="Knowledge Base ID"
        ),
        sa.Column("file_tag_id", sa.Text(), nullable=True, comment="File Tag ID"),
        sa.Column("file_id", sa.String(length=36), nullable=False, comment="File ID"),
        sa.Column(
            "file_key", sa.String(length=255), nullable=False, comment="File oss key"
        ),
        sa.Column(
            "file_name", sa.String(length=255), nullable=True, comment="File name"
        ),
        sa.Column("file_text", sa.Text(), nullable=True, comment="File text"),
        sa.Column("meta_data", sa.Text(), nullable=True, comment="Meta Data"),
        sa.Column("extra_data", sa.Text(), nullable=True, comment="Extra Data"),
        sa.Column(
            "created_user_id",
            sa.String(length=36),
            nullable=True,
            comment="created user ID",
        ),
        sa.Column(
            "updated_user_id",
            sa.String(length=36),
            nullable=True,
            comment="updated user ID",
        ),
        sa.Column("created", sa.TIMESTAMP(), nullable=False, comment="Creation time"),
        sa.Column("updated", sa.TIMESTAMP(), nullable=False, comment="Update time"),
        sa.PrimaryKeyConstraint("id"),
    )
    with op.batch_alter_table("knowledge_file", schema=None) as batch_op:
        batch_op.create_index(
            batch_op.f("ix_knowledge_file_file_id"), ["file_id"], unique=False
        )
        batch_op.create_index(
            batch_op.f("ix_knowledge_file_kb_id"), ["kb_id"], unique=False
        )

    op.create_table(
        "tag",
        sa.Column("id", mysql.BIGINT(), autoincrement=True, nullable=False),
        sa.Column("tag_id", sa.String(length=36), nullable=False, comment="Tag ID"),
        sa.Column(
            "tag_domain", sa.String(length=64), nullable=False, comment="Tag domain"
        ),
        sa.Column("tag_type", sa.String(length=64), nullable=False, comment="Tag type"),
        sa.Column(
            "tag_name", sa.String(length=255), nullable=False, comment="Tag name"
        ),
        sa.Column("meta_data", sa.Text(), nullable=True, comment="Meta Data"),
        sa.Column("extra_data", sa.Text(), nullable=True, comment="Extra Data"),
        sa.Column(
            "created_user_id",
            sa.String(length=36),
            nullable=True,
            comment="created user ID",
        ),
        sa.Column(
            "updated_user_id",
            sa.String(length=36),
            nullable=True,
            comment="updated user ID",
        ),
        sa.Column("created", sa.TIMESTAMP(), nullable=False, comment="Creation time"),
        sa.Column("updated", sa.TIMESTAMP(), nullable=False, comment="Update time"),
        sa.PrimaryKeyConstraint("id"),
    )
    with op.batch_alter_table("tag", schema=None) as batch_op:
        batch_op.create_index(batch_op.f("ix_tag_tag_id"), ["tag_id"], unique=False)

    # ### end Alembic commands ###


def downgrade():
    with op.batch_alter_table("tag", schema=None) as batch_op:
        batch_op.drop_index(batch_op.f("ix_tag_tag_id"))

    op.drop_table("tag")
    with op.batch_alter_table("knowledge_file", schema=None) as batch_op:
        batch_op.drop_index(batch_op.f("ix_knowledge_file_kb_id"))
        batch_op.drop_index(batch_op.f("ix_knowledge_file_file_id"))

    op.drop_table("knowledge_file")
    with op.batch_alter_table("knowledge_chunk", schema=None) as batch_op:
        batch_op.drop_index(batch_op.f("ix_knowledge_chunk_kb_id"))
        batch_op.drop_index(batch_op.f("ix_knowledge_chunk_file_id"))
        batch_op.drop_index(batch_op.f("ix_knowledge_chunk_chunk_id"))

    op.drop_table("knowledge_chunk")
    with op.batch_alter_table("knowledge_base", schema=None) as batch_op:
        batch_op.drop_index(batch_op.f("ix_knowledge_base_kb_id"))

    op.drop_table("knowledge_base")
    # ### end Alembic commands ###
