"""Add knowledge table

Revision ID: 6a39f3d8e55c
Revises: c0fbf31ca0db
Create Date: 2024-10-01 14:02:35.241684

"""

import json

import sqlalchemy as sa
from alembic import op
from sqlalchemy.sql import column, select, table

revision = '6a39f3d8e55c'
down_revision = 'c0fbf31ca0db'
branch_labels = None
depends_on = None


def upgrade():
    conn = op.get_bind()
    inspector = sa.inspect(conn)
    existing_tables = set(inspector.get_table_names())

    if 'knowledge' not in existing_tables:
        # Creating the 'knowledge' table
        print('Creating knowledge table')
        knowledge_table = op.create_table(
            'knowledge',
            sa.Column('id', sa.Text(), primary_key=True),
            sa.Column('user_id', sa.Text(), nullable=False),
            sa.Column('name', sa.Text(), nullable=False),
            sa.Column('description', sa.Text(), nullable=True),
            sa.Column('data', sa.JSON(), nullable=True),
            sa.Column('meta', sa.JSON(), nullable=True),
            sa.Column('created_at', sa.BigInteger(), nullable=False),
            sa.Column('updated_at', sa.BigInteger(), nullable=True),
        )

        print('Migrating data from document table to knowledge table')
        # Representation of the existing 'document' table
        document_table = table(
            'document',
            column('collection_name', sa.String()),
            column('user_id', sa.String()),
            column('name', sa.String()),
            column('title', sa.Text()),
            column('content', sa.Text()),
            column('timestamp', sa.BigInteger()),
        )

        # Select all from existing document table
        documents = conn.execute(
            select(
                document_table.c.collection_name,
                document_table.c.user_id,
                document_table.c.name,
                document_table.c.title,
                document_table.c.content,
                document_table.c.timestamp,
            )
        )

        # Insert data into knowledge table from document table
        for doc in documents:
            conn.execute(
                knowledge_table.insert().values(
                    id=doc.collection_name,
                    user_id=doc.user_id,
                    description=doc.name,
                    meta={
                        'legacy': True,
                        'document': True,
                        'tags': json.loads(doc.content or '{}').get('tags', []),
                    },
                    name=doc.title,
                    created_at=doc.timestamp,
                    updated_at=doc.timestamp,
                )
            )


def downgrade():
    op.drop_table('knowledge')
