#![allow(unused_imports, dead_code)]

pub mod common;

pub use common::{TestContext, features::*, setup::*};
use pretty_assertions::assert_eq;
use sea_orm::{DatabaseConnection, DerivePartialModel, entity::prelude::*, entity::*};
use serde_json::json;

#[cfg(feature = "postgres-vector")]
mod test {
    use pretty_assertions::assert_eq;

    use super::*;

    #[sea_orm_macros::test]
    fn main() -> Result<(), DbErr> {
        let ctx = TestContext::new("embedding_tests");
        create_embedding_table(&ctx.db)?;
        insert_embedding(&ctx.db)?;
        update_embedding(&ctx.db)?;
        select_embedding(&ctx.db)?;
        ctx.delete();

        Ok(())
    }

    pub fn insert_embedding(db: &DatabaseConnection) -> Result<(), DbErr> {
        use embedding::*;

        assert_eq!(
            Model {
                id: 1,
                embedding: PgVector::from(vec![1.]),
            }
            .into_active_model()
            .insert(db)?,
            Model {
                id: 1,
                embedding: PgVector::from(vec![1.]),
            }
        );

        assert_eq!(
            Model {
                id: 2,
                embedding: PgVector::from(vec![1., 2.]),
            }
            .into_active_model()
            .insert(db)?,
            Model {
                id: 2,
                embedding: PgVector::from(vec![1., 2.]),
            }
        );

        assert_eq!(
            Model {
                id: 3,
                embedding: PgVector::from(vec![1., 2., 3.]),
            }
            .into_active_model()
            .insert(db)?,
            Model {
                id: 3,
                embedding: PgVector::from(vec![1., 2., 3.]),
            }
        );

        assert_eq!(
            Entity::find_by_id(3).into_json().one(db)?,
            Some(json!({
                "id": 3,
                "embedding": [1., 2., 3.],
            }))
        );

        Ok(())
    }

    pub fn update_embedding(db: &DatabaseConnection) -> Result<(), DbErr> {
        use embedding::*;

        let model = Entity::find_by_id(1).one(db)?.unwrap();

        ActiveModel {
            embedding: Set(PgVector::from(vec![10.])),
            ..model.into_active_model()
        }
        .update(db)?;

        ActiveModel {
            id: Unchanged(3),
            embedding: Set(PgVector::from(vec![10., 20., 30.])),
        }
        .update(db)?;

        Ok(())
    }

    pub fn select_embedding(db: &DatabaseConnection) -> Result<(), DbErr> {
        use embedding::*;

        #[derive(DerivePartialModel, Debug, PartialEq)]
        #[sea_orm(entity = "Entity")]
        struct PartialSelectResult {
            embedding: PgVector,
        }

        let result = Entity::find_by_id(1)
            .into_partial_model::<PartialSelectResult>()
            .one(db)?;

        assert_eq!(
            result,
            Some(PartialSelectResult {
                embedding: PgVector::from(vec![10.]),
            })
        );

        let result = Entity::find_by_id(2)
            .into_partial_model::<PartialSelectResult>()
            .one(db)?;

        assert_eq!(
            result,
            Some(PartialSelectResult {
                embedding: PgVector::from(vec![1., 2.]),
            })
        );

        let result = Entity::find_by_id(3)
            .into_partial_model::<PartialSelectResult>()
            .one(db)?;

        assert_eq!(
            result,
            Some(PartialSelectResult {
                embedding: PgVector::from(vec![10., 20., 30.]),
            })
        );

        Ok(())
    }
}
