orign 0.2.3

// src/query.rs
use crate::entities::buffer;
use crate::entities::feedback;
use crate::entities::human;
use crate::entities::llm;
use crate::entities::model_deployment;
use crate::entities::training_job;
use sea_orm::*;
use sea_orm::{ColumnTrait, DatabaseConnection, DbErr, EntityTrait, QueryFilter};

pub struct Query;

impl Query {
    pub async fn find_deployment_by_id(
        db: &DatabaseConnection,
        id: &str,
    ) -> Result<Option<model_deployment::Model>, DbErr> {
        model_deployment::Entity::find_by_id(id).one(db).await
    }

    pub async fn find_deployments_by_owner(
        db: &DatabaseConnection,
        owner_id: &str,
    ) -> Result<Vec<model_deployment::Model>, DbErr> {
        model_deployment::Entity::find()
            .filter(model_deployment::Column::OwnerId.eq(owner_id))
            .all(db)
            .await
    }

    pub async fn find_deployments_by_framework(
        db: &DatabaseConnection,
        framework: &str,
    ) -> Result<Vec<model_deployment::Model>, DbErr> {
        model_deployment::Entity::find()
            .filter(model_deployment::Column::Framework.eq(framework))
            .all(db)
            .await
    }

    pub async fn find_all_deployments(
        db: &DatabaseConnection,
    ) -> Result<Vec<model_deployment::Model>, DbErr> {
        model_deployment::Entity::find().all(db).await
    }

    pub async fn find_deployments_by_kind(
        db: &DatabaseConnection,
        kind: &str,
    ) -> Result<Vec<model_deployment::Model>, DbErr> {
        model_deployment::Entity::find()
            .filter(model_deployment::Column::Kind.eq(kind))
            .all(db)
            .await
    }

    // Helper function to find matching deployments based on model and provider criteria
    pub async fn find_matching_deployments(
        db: &DatabaseConnection,
        model: Option<&str>,
        framework: Option<&str>,
        kind: &str, // kind is now required
    ) -> Result<Vec<model_deployment::Model>, DbErr> {
        // Start with kind filter as it's always required
        let mut query =
            model_deployment::Entity::find().filter(model_deployment::Column::Kind.eq(kind));

        // Add provider filter if specified
        if let Some(framework) = framework {
            query = query.filter(model_deployment::Column::Framework.eq(framework));
        }

        // Execute the query
        let deployments = query.all(db).await?;

        // If model is specified, filter deployments by checking model field and params
        if let Some(model_name) = model {
            Ok(deployments
                .into_iter()
                .filter(|deployment| {
                    // Check the model field first
                    if deployment.model == model_name {
                        return true;
                    }

                    // Then check the params if they exist
                    if let Some(params) = &deployment.params {
                        if let Some(param_model) = params.get("model") {
                            if let Some(model_str) = param_model.as_str() {
                                return model_str == model_name;
                            }
                        }
                    }
                    false
                })
                .collect())
        } else {
            Ok(deployments)
        }
    }

    pub async fn find_training_job_by_id(
        db: &DatabaseConnection,
        id: &str,
    ) -> Result<Option<training_job::Model>, DbErr> {
        training_job::Entity::find_by_id(id).one(db).await
    }

    pub async fn find_training_jobs_by_owners(
        db: &DatabaseConnection,
        owner_ids: &[&str],
    ) -> Result<Vec<training_job::Model>, DbErr> {
        training_job::Entity::find()
            .filter(training_job::Column::OwnerId.is_in(owner_ids.iter().copied()))
            .all(db)
            .await
    }

    pub async fn find_training_jobs_by_queue(
        db: &DatabaseConnection,
        queue_name: &str,
    ) -> Result<Vec<training_job::Model>, DbErr> {
        training_job::Entity::find()
            .filter(training_job::Column::Queue.eq(queue_name))
            .all(db)
            .await
    }

    pub async fn find_training_jobs_by_owner(
        db: &DatabaseConnection,
        owner_id: &str,
    ) -> Result<Vec<training_job::Model>, DbErr> {
        training_job::Entity::find()
            .filter(training_job::Column::OwnerId.eq(owner_id))
            .all(db)
            .await
    }

    // pub async fn find_training_job_by_id_and_owners(
    //     db: &DatabaseConnection,
    //     id: &str,
    //     owner_ids: &[&str],
    // ) -> Result<training_job::Model, DbErr> {
    //     training_job::Entity::find()
    //         .filter(
    //             training_job::Column::Id
    //                 .eq(id.to_owned())
    //                 .and(training_job::Column::OwnerId.is_in(owner_ids.iter().copied())),
    //         )
    //         .one(db)
    //         .await?
    //         .ok_or_else(|| DbErr::RecordNotFound(format!("No training job found with id {}", id)))
    // }

    // pub async fn find_training_job_by_id_and_owner(
    //     db: &DatabaseConnection,
    //     id: &str,
    //     owner_id: &str,
    // ) -> Result<training_job::Model, DbErr> {
    //     training_job::Entity::find()
    //         .filter(
    //             training_job::Column::Id
    //                 .eq(id.to_owned())
    //                 .and(training_job::Column::OwnerId.eq(owner_id.to_owned())),
    //         )
    //         .one(db)
    //         .await?
    //         .ok_or_else(|| DbErr::RecordNotFound(format!("No training job found with id {}", id)))
    // }

    pub async fn find_training_job_by_resource_name(
        db: &DatabaseConnection,
        resource_name: &str,
    ) -> Result<Option<training_job::Model>, DbErr> {
        training_job::Entity::find()
            .filter(training_job::Column::ResourceName.eq(resource_name))
            .one(db)
            .await
    }

    pub async fn find_buffer_by_id(
        db: &DatabaseConnection,
        id: &str,
    ) -> Result<Option<buffer::Model>, DbErr> {
        buffer::Entity::find_by_id(id).one(db).await
    }

    pub async fn find_buffer_by_name_and_owners(
        db: &DatabaseConnection,
        name: &str,
        namespace: &str,
        owner_ids: &[&str],
    ) -> Result<Option<buffer::Model>, DbErr> {
        buffer::Entity::find()
            .filter(buffer::Column::Name.eq(name))
            .filter(buffer::Column::Namespace.eq(namespace))
            .filter(buffer::Column::OwnerId.is_in(owner_ids.iter().copied()))
            .one(db)
            .await
    }

    pub async fn find_buffer_by_name(
        db: &DatabaseConnection,
        namespace: &str,
        name: &str,
    ) -> Result<Option<buffer::Model>, DbErr> {
        buffer::Entity::find()
            .filter(buffer::Column::Name.eq(name))
            .filter(buffer::Column::Namespace.eq(namespace))
            .one(db)
            .await
    }

    pub async fn find_buffers_by_owner(
        db: &DatabaseConnection,
        owner_id: &str,
    ) -> Result<Vec<buffer::Model>, DbErr> {
        buffer::Entity::find()
            .filter(buffer::Column::OwnerId.eq(owner_id))
            .all(db)
            .await
    }

    pub async fn find_buffers_by_owners(
        db: &DatabaseConnection,
        owner_ids: &[&str],
    ) -> Result<Vec<buffer::Model>, DbErr> {
        buffer::Entity::find()
            .filter(buffer::Column::OwnerId.is_in(owner_ids.iter().cloned()))
            .all(db)
            .await
    }

    pub async fn find_buffer_by_id_and_owners(
        db: &DatabaseConnection,
        id: &str,
        owner_ids: &[&str],
    ) -> Result<Option<buffer::Model>, DbErr> {
        buffer::Entity::find()
            .filter(buffer::Column::Id.eq(id))
            .filter(buffer::Column::OwnerId.is_in(owner_ids.iter().copied()))
            .one(db)
            .await
    }

    pub async fn find_llms_by_owners(
        db: &DatabaseConnection,
        owner_ids: &[&str],
    ) -> Result<Vec<llm::Model>, DbErr> {
        llm::Entity::find()
            .filter(llm::Column::OwnerId.is_in(owner_ids.iter().copied()))
            .all(db)
            .await
    }

    pub async fn find_llm_by_name_and_namespace(
        db: &DatabaseConnection,
        name: &str,
        namespace: &str,
    ) -> Result<Option<llm::Model>, DbErr> {
        llm::Entity::find()
            .filter(llm::Column::Name.eq(name))
            .filter(llm::Column::Namespace.eq(namespace))
            .one(db)
            .await
    }

    pub async fn find_llm_by_name_and_namespace_and_owners(
        db: &DatabaseConnection,
        name: &str,
        namespace: &str,
        owner_ids: &[&str],
    ) -> Result<Option<llm::Model>, DbErr> {
        llm::Entity::find()
            .filter(llm::Column::Name.eq(name))
            .filter(llm::Column::Namespace.eq(namespace))
            .filter(llm::Column::OwnerId.is_in(owner_ids.iter().copied()))
            .one(db)
            .await
    }

    pub async fn find_humans_by_owners(
        db: &DatabaseConnection,
        owner_ids: &[&str],
    ) -> Result<Vec<human::Model>, DbErr> {
        human::Entity::find()
            .filter(human::Column::OwnerId.is_in(owner_ids.iter().copied()))
            .all(db)
            .await
    }

    pub async fn find_human_by_name_and_namespace_and_owners(
        db: &DatabaseConnection,
        name: &str,
        namespace: &str,
        owner_ids: &[&str],
    ) -> Result<Option<human::Model>, DbErr> {
        human::Entity::find()
            .filter(human::Column::Name.eq(name))
            .filter(human::Column::Namespace.eq(namespace))
            .filter(human::Column::OwnerId.is_in(owner_ids.iter().copied()))
            .one(db)
            .await
    }
    pub async fn find_feedback_by_id_and_human_id(
        db: &DatabaseConnection,
        feedback_id: &str,
        human_id: &str,
    ) -> Result<Option<feedback::Model>, DbErr> {
        feedback::Entity::find_by_id(feedback_id)
            .filter(feedback::Column::HumanId.eq(human_id))
            .one(db)
            .await
    }
}