vantage-surrealdb 0.5.9

use serde_json::Value;
use std::marker::PhantomData;
use std::ops::Deref;
use vantage_core::{Context, Result, vantage_error};
use vantage_expressions::protocol::result::{self, QueryResult};
use vantage_expressions::{AssociatedQueryable, Expression, Order, Selectable};
use vantage_table::Entity;

use crate::select::SurrealSelect;
use crate::surreal_return::SurrealReturn;
use crate::{SurrealDB, protocol::SurrealQueriable};

/// SurrealDB-specific trait for associated queries that can be executed
/// This trait extends AssociatedQueryable with Value conversion
#[async_trait::async_trait]
pub trait SurrealAssociatedQueryable<R>: AssociatedQueryable<R> {
    /// Execute the query and return raw Value
    async fn get_as_value(&self) -> Result<Value>;
}

/// SurrealDB-specific associated query that combines any SurrealQueriable with SurrealDB datasource
pub struct SurrealAssociated<Q: SurrealQueriable, R> {
    pub query: Q,
    pub datasource: SurrealDB,
    _result: PhantomData<R>,
}

impl<Q: SurrealQueriable + Clone, R> Clone for SurrealAssociated<Q, R> {
    fn clone(&self) -> Self {
        Self {
            query: self.query.clone(),
            datasource: self.datasource.clone(),
            _result: PhantomData,
        }
    }
}

impl<Q: SurrealQueriable + std::fmt::Debug, R> std::fmt::Debug for SurrealAssociated<Q, R> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("SurrealAssociated")
            .field("query", &self.query)
            .field("datasource", &"SurrealDB{...}")
            .finish()
    }
}

impl<Q: SurrealQueriable, R> SurrealAssociated<Q, R> {
    pub fn new(query: Q, datasource: SurrealDB) -> Self {
        Self {
            query,
            datasource,
            _result: PhantomData,
        }
    }
}

// Wrapper methods for SurrealSelect that preserve the SurrealAssociated wrapper
impl<T: QueryResult, R> SurrealAssociated<SurrealSelect<T>, R> {
    pub fn with_limit(mut self, limit: Option<i64>, skip: Option<i64>) -> Self {
        self.query.set_limit(limit, skip);
        self
    }

    pub fn with_order_by(
        mut self,
        expression: vantage_expressions::Expression,
        order: Order,
    ) -> Self {
        self.query.add_order_by(expression, order);
        self
    }

    pub fn with_field(mut self, field: impl Into<String>) -> Self {
        self.query.add_field(field);
        self
    }

    pub fn with_condition(mut self, condition: vantage_expressions::Expression) -> Self {
        self.query.add_where_condition(condition);
        self
    }
}

// Preview method implementation for SurrealSelect types
impl<T: QueryResult, R> SurrealAssociated<SurrealSelect<T>, R> {
    pub fn preview(&self) -> String {
        self.query.preview()
    }
}

impl<Q: SurrealQueriable, R> Deref for SurrealAssociated<Q, R> {
    type Target = Q;

    fn deref(&self) -> &Self::Target {
        &self.query
    }
}

impl<Q: SurrealQueriable + Selectable + Clone, R: Send + Sync> Selectable
    for SurrealAssociated<Q, R>
{
    fn add_source(&mut self, source: impl Into<vantage_expressions::Expr>, alias: Option<String>) {
        self.query.add_source(source, alias)
    }

    fn add_field(&mut self, field: impl Into<String>) {
        self.query.add_field(field)
    }

    fn add_expression(&mut self, expression: vantage_expressions::Expression) {
        self.query.add_expression(expression)
    }

    fn add_where_condition(&mut self, condition: vantage_expressions::Expression) {
        self.query.add_where_condition(condition)
    }

    fn set_distinct(&mut self, distinct: bool) {
        self.query.set_distinct(distinct)
    }

    fn add_order_by(&mut self, expression: Expression, order: Order) {
        self.query.add_order_by(expression, order)
    }

    fn add_group_by(&mut self, expression: vantage_expressions::Expression) {
        self.query.add_group_by(expression)
    }

    fn set_limit(&mut self, limit: Option<i64>, skip: Option<i64>) {
        self.query.set_limit(limit, skip)
    }

    fn clear_fields(&mut self) {
        self.query.clear_fields()
    }

    fn clear_where_conditions(&mut self) {
        self.query.clear_where_conditions()
    }

    fn clear_order_by(&mut self) {
        self.query.clear_order_by()
    }

    fn clear_group_by(&mut self) {
        self.query.clear_group_by()
    }

    fn has_fields(&self) -> bool {
        self.query.has_fields()
    }

    fn has_where_conditions(&self) -> bool {
        self.query.has_where_conditions()
    }

    fn has_order_by(&self) -> bool {
        self.query.has_order_by()
    }

    fn has_group_by(&self) -> bool {
        self.query.has_group_by()
    }

    fn is_distinct(&self) -> bool {
        self.query.is_distinct()
    }

    fn get_limit(&self) -> Option<i64> {
        self.query.get_limit()
    }

    fn get_skip(&self) -> Option<i64> {
        self.query.get_skip()
    }

    fn as_field(&self, field: impl Into<String>) -> vantage_expressions::Expression {
        self.query.as_field(field)
    }

    fn as_count(&self) -> vantage_expressions::Expression {
        self.query.as_count()
    }

    fn as_sum(&self, column: vantage_expressions::Expression) -> vantage_expressions::Expression {
        self.query.as_sum(column)
    }
}

impl<Q: SurrealQueriable + Into<Expression>, R> From<SurrealAssociated<Q, R>> for Expression {
    fn from(val: SurrealAssociated<Q, R>) -> Self {
        val.query.into()
    }
}

// Implementation for SurrealSelect<result::Rows> - returns Vec<Entity>
#[async_trait::async_trait]
impl<E: Entity> AssociatedQueryable<Vec<E>>
    for SurrealAssociated<SurrealSelect<result::Rows>, Vec<E>>
{
    async fn get(&self) -> Result<Vec<E>> {
        let raw_result = self.query.get(&self.datasource).await;
        let entities = raw_result
            .into_iter()
            .map(|item| serde_json::from_value(Value::Object(item)))
            .collect::<std::result::Result<Vec<E>, _>>()
            .context("Couldn't convert from value")?;
        Ok(entities)
    }
}

#[async_trait::async_trait]
impl<E: Entity> SurrealAssociatedQueryable<Vec<E>>
    for SurrealAssociated<SurrealSelect<result::Rows>, Vec<E>>
{
    async fn get_as_value(&self) -> Result<Value> {
        let raw_result = self.query.get(&self.datasource).await;
        let json_value = Value::Array(raw_result.into_iter().map(Value::Object).collect());
        Ok(json_value)
    }
}

// Implementation for SurrealSelect<result::SingleRow> - returns Entity
#[async_trait::async_trait]
impl<E: Entity> AssociatedQueryable<E> for SurrealAssociated<SurrealSelect<result::SingleRow>, E> {
    async fn get(&self) -> Result<E> {
        let raw_result = self.query.get(&self.datasource).await;
        let entity: E = serde_json::from_value(Value::Object(raw_result))
            .context("Couldn't convert from value")?;
        Ok(entity)
    }
}

#[async_trait::async_trait]
impl<E: Entity> SurrealAssociatedQueryable<E>
    for SurrealAssociated<SurrealSelect<result::SingleRow>, E>
{
    async fn get_as_value(&self) -> Result<Value> {
        let raw_result = self.query.get(&self.datasource).await;
        Ok(Value::Object(raw_result))
    }
}

// Implementation for SurrealSelect<result::List> - returns Vec<Value>
#[async_trait::async_trait]
impl AssociatedQueryable<Vec<Value>>
    for SurrealAssociated<SurrealSelect<result::List>, Vec<Value>>
{
    async fn get(&self) -> Result<Vec<Value>> {
        let raw_result = self.query.get(&self.datasource).await;
        Ok(raw_result)
    }
}

#[async_trait::async_trait]
impl SurrealAssociatedQueryable<Vec<Value>>
    for SurrealAssociated<SurrealSelect<result::List>, Vec<Value>>
{
    async fn get_as_value(&self) -> Result<Value> {
        let raw_result = self.query.get(&self.datasource).await;
        Ok(Value::Array(raw_result))
    }
}

// Implementation for SurrealSelect<result::Single> - returns Value
#[async_trait::async_trait]
impl AssociatedQueryable<Value> for SurrealAssociated<SurrealSelect<result::Single>, Value> {
    async fn get(&self) -> Result<Value> {
        let raw_result = self.query.get(&self.datasource).await;
        Ok(raw_result)
    }
}

#[async_trait::async_trait]
impl SurrealAssociatedQueryable<Value> for SurrealAssociated<SurrealSelect<result::Single>, Value> {
    async fn get_as_value(&self) -> Result<Value> {
        let raw_result = self.query.get(&self.datasource).await;
        Ok(raw_result)
    }
}

// Implementation for SurrealReturn - returns Value
#[async_trait::async_trait]
impl AssociatedQueryable<Value> for SurrealAssociated<SurrealReturn, Value> {
    async fn get(&self) -> Result<Value> {
        let raw_result = self.query.get(&self.datasource).await;
        Ok(raw_result)
    }
}

#[async_trait::async_trait]
impl SurrealAssociatedQueryable<Value> for SurrealAssociated<SurrealReturn, Value> {
    async fn get_as_value(&self) -> Result<Value> {
        let raw_result = self.query.get(&self.datasource).await;
        Ok(raw_result)
    }
}

// Implementation for SurrealReturn - returns i64 (for count queries)
#[async_trait::async_trait]
impl AssociatedQueryable<i64> for SurrealAssociated<SurrealReturn, i64> {
    async fn get(&self) -> Result<i64> {
        let raw_result = self.query.get(&self.datasource).await;
        // SurrealDB count returns [number], so extract the first element
        if let Value::Array(ref arr) = raw_result
            && let Some(Value::Number(num)) = arr.first()
            && let Some(count) = num.as_i64()
        {
            return Ok(count);
        }
        Err(vantage_error!("Failed to parse count result as i64"))
    }
}

#[async_trait::async_trait]
impl SurrealAssociatedQueryable<i64> for SurrealAssociated<SurrealReturn, i64> {
    async fn get_as_value(&self) -> Result<Value> {
        let raw_result = self.query.get(&self.datasource).await;
        Ok(raw_result)
    }
}