use sea_orm_migration::prelude::*;

#[derive(DeriveMigrationName)]
pub struct Migration;

const AI_CHUNKS_TABLE: &str = "ai_chunks";
const VECTOR_INDEX_NAME: &str = "idx_ai_chunks_embedding_hnsw";
const EMBEDDING_DIMENSION: i32 = 384;

#[async_trait::async_trait]
impl MigrationTrait for Migration {
    async fn up(&self, manager: &SchemaManager) -> Result<(), DbErr> {
        manager
            .get_connection()
            .execute_unprepared("CREATE EXTENSION IF NOT EXISTS vector")
            .await?;

        if manager.has_column(AI_CHUNKS_TABLE, "embedding").await? {
            manager
                .get_connection()
                .execute_unprepared("ALTER TABLE ai_chunks DROP COLUMN embedding")
                .await?;
        }

        if !manager.has_column(AI_CHUNKS_TABLE, "embedding").await? {
            manager
                .get_connection()
                .execute_unprepared(&format!(
                    "ALTER TABLE ai_chunks ADD COLUMN embedding vector({EMBEDDING_DIMENSION})"
                ))
                .await?;
        }

        manager
            .get_connection()
            .execute_unprepared("TRUNCATE TABLE ai_chunks RESTART IDENTITY")
            .await?;

        manager
            .get_connection()
            .execute_unprepared(&format!(
                "CREATE INDEX IF NOT EXISTS {VECTOR_INDEX_NAME} ON ai_chunks USING hnsw (embedding vector_cosine_ops)"
            ))
            .await?;

        Ok(())
    }

    async fn down(&self, manager: &SchemaManager) -> Result<(), DbErr> {
        manager
            .get_connection()
            .execute_unprepared(&format!("DROP INDEX IF EXISTS {VECTOR_INDEX_NAME}"))
            .await?;

        if manager.has_column(AI_CHUNKS_TABLE, "embedding").await? {
            manager
                .get_connection()
                .execute_unprepared("ALTER TABLE ai_chunks DROP COLUMN embedding")
                .await?;
        }

        manager
            .get_connection()
            .execute_unprepared("ALTER TABLE ai_chunks ADD COLUMN embedding jsonb")
            .await?;

        Ok(())
    }
}
