"""init

Revision ID: 955beece0126
Revises:
Create Date: 2024-12-23 14:46:07.728126

"""

from typing import Sequence, Union

import sqlalchemy as sa
import sqlite_vec
from alembic import op
from sqlmodel.sql.sqltypes import AutoString

# revision identifiers, used by Alembic.
revision: str = "955beece0126"
down_revision: Union[str, None] = None
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:

    # Add debug logging
    import logging

    logging.debug("Starting database upgrade...")

    # Create tables first
    logging.debug("Creating tables...")
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table(
        "contextmessageset",
        sa.Column("id", sa.Integer(), nullable=False),
        sa.Column("created_at", sa.DateTime(), nullable=False),
        sa.Column("updated_at", sa.DateTime(), nullable=False),
        sa.Column("user_id", sa.Integer(), nullable=False),
        sa.Column("message_ids", sa.Text(), nullable=True),
        sa.Column("is_active", sa.Boolean(), nullable=True),
        sa.PrimaryKeyConstraint("id"),
        sa.UniqueConstraint("user_id", "is_active"),
    )
    op.create_index(op.f("ix_contextmessageset_id"), "contextmessageset", ["id"], unique=False)
    op.create_table(
        "goal",
        sa.Column("id", sa.Integer(), nullable=False),
        sa.Column("created_at", sa.DateTime(), nullable=False),
        sa.Column("updated_at", sa.DateTime(), nullable=False),
        sa.Column("user_id", sa.Integer(), nullable=False),
        sa.Column("name", AutoString(), nullable=False),
        sa.Column("status_updates", sa.Text(), server_default="[]", nullable=False),
        sa.Column("description", AutoString(), nullable=True),
        sa.Column("strategy", AutoString(), nullable=True),
        sa.Column("end_condition", AutoString(), nullable=True),
        sa.Column("is_active", sa.Boolean(), nullable=True),
        sa.Column("priority", sa.Integer(), nullable=True),
        sa.Column("target_completion_time", sa.DateTime(), nullable=True),
        sa.PrimaryKeyConstraint("id"),
        sa.UniqueConstraint("user_id", "name", "is_active"),
    )
    op.create_index(op.f("ix_goal_id"), "goal", ["id"], unique=False)
    op.create_table(
        "memory",
        sa.Column("id", sa.Integer(), nullable=False),
        sa.Column("created_at", sa.DateTime(), nullable=False),
        sa.Column("updated_at", sa.DateTime(), nullable=False),
        sa.Column("user_id", sa.Integer(), nullable=False),
        sa.Column("name", AutoString(), nullable=False),
        sa.Column("text", AutoString(), nullable=False),
        sa.Column("is_active", sa.Boolean(), nullable=True),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_index(op.f("ix_memory_id"), "memory", ["id"], unique=False)
    op.create_table(
        "message",
        sa.Column("id", sa.Integer(), nullable=False),
        sa.Column("created_at", sa.DateTime(), nullable=False),
        sa.Column("updated_at", sa.DateTime(), nullable=False),
        sa.Column("user_id", sa.Integer(), nullable=False),
        sa.Column("role", AutoString(), nullable=False),
        sa.Column("content", AutoString(), nullable=True),
        sa.Column("model", AutoString(), nullable=True),
        sa.Column("tool_calls", sa.Text(), nullable=True),
        sa.Column("tool_call_id", AutoString(), nullable=True),
        sa.Column("memory_metadata", sa.Text(), nullable=True),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_index(op.f("ix_message_id"), "message", ["id"], unique=False)
    op.create_table(
        "user",
        sa.Column("id", sa.Integer(), nullable=False),
        sa.Column("token", AutoString(), nullable=False),
        sa.Column("created_at", sa.DateTime(), nullable=False),
        sa.Column("updated_at", sa.DateTime(), nullable=False),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_index(op.f("ix_user_id"), "user", ["id"], unique=False)
    op.create_table(
        "userpreference",
        sa.Column("id", sa.Integer(), nullable=False),
        sa.Column("created_at", sa.DateTime(), nullable=False),
        sa.Column("updated_at", sa.DateTime(), nullable=False),
        sa.Column("user_id", sa.Integer(), nullable=False),
        sa.Column("preferred_name", AutoString(), nullable=True),
        sa.Column("system_persona", AutoString(), nullable=True),
        sa.Column("full_name", AutoString(), nullable=True),
        sa.Column("is_active", sa.Boolean(), nullable=True),
        sa.PrimaryKeyConstraint("id"),
        sa.UniqueConstraint("user_id", "is_active"),
    )
    op.create_index(op.f("ix_userpreference_id"), "userpreference", ["id"], unique=False)

    # Then handle vector extension
    logging.debug("Attempting to load sqlite_vec extension...")

    connection = op.get_bind().connection

    try:

        connection.enable_load_extension(True)  # type: ignore
        sqlite_vec.load(connection)  # type: ignore
        connection.enable_load_extension(False)  # type: ignore
    except Exception as e:
        logging.error(f"Failed to load sqlite_vec extension: {e}")
        raise

    logging.debug("Attempting to create vectorstorage table...")

    op.execute(
        """
        CREATE VIRTUAL TABLE IF NOT EXISTS vectorstorage USING vec0(
            id INTEGER PRIMARY KEY,  -- This aliases rowid
            source_type TEXT,
            source_id INTEGER,
            embedding_data FLOAT[1536],
            embedding_text_md5 TEXT
        )
    """
    )

    logging.debug("Database upgrade completed successfully")
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.execute("DROP TABLE IF EXISTS vectorstorage")
    op.drop_index(op.f("ix_userpreference_id"), table_name="userpreference")
    op.drop_table("userpreference")
    op.drop_index(op.f("ix_user_id"), table_name="user")
    op.drop_table("user")
    op.drop_index(op.f("ix_message_id"), table_name="message")
    op.drop_table("message")
    op.drop_index(op.f("ix_memory_id"), table_name="memory")
    op.drop_table("memory")
    op.drop_index(op.f("ix_goal_id"), table_name="goal")
    op.drop_table("goal")
    op.drop_index(op.f("ix_contextmessageset_id"), table_name="contextmessageset")
    op.drop_table("contextmessageset")
    # ### end Alembic commands ###
