// Copyright 2024 The NativeLink Authors. All rights reserved.

//

// Licensed under the Apache License, Version 2.0 (the "License");

// you may not use this file except in compliance with the License.

// You may obtain a copy of the License at

//

//    http://www.apache.org/licenses/LICENSE-2.0

//

// Unless required by applicable law or agreed to in writing, software

// distributed under the License is distributed on an "AS IS" BASIS,

// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

// See the License for the specific language governing permissions and

// limitations under the License.

use std::collections::VecDeque;

use fred::error::{Error as RedisError, ErrorKind as RedisErrorKind};

use fred::interfaces::RediSearchInterface;

use fred::types::redisearch::FtAggregateOptions;

use fred::types::{FromValue, Map as RedisMap, Value as RedisValue};

use futures::Stream;

/// Calls `FT_AGGREGATE` in redis. Fred does not properly support this command

/// so we have to manually handle it.

pub(crate) async fn ft_aggregate<C, I, Q>(

    client: C,

    index: I,

    query: Q,

    options: FtAggregateOptions,

) -> Result<impl Stream<Item = Result<RedisMap, RedisError>> + Send, RedisError>

where

    C: RediSearchInterface,

    I: Into<bytes_utils::string::Str>,

    Q: Into<bytes_utils::string::Str>,

{

    struct State<C: RediSearchInterface> {

        client: C,

        index: bytes_utils::string::Str,

        data: RedisCursorData,

    }

    let index = index.into();

    let query = query.into();

    let 
data21
: 
RedisCursorData21
 = client.ft_aggregate(index.clone(), query, options).await
?1
;

    let state = State {

        client,

        index,

        data,

    };

    Ok(futures::stream::unfold(

        Some(state),

        move |maybe_state| async move {

            let mut state = maybe_state
?0
;

            loop {

                if let Some(
map7
) = state.data.data.pop_front() {

  Branch (56:24): [True: 0, False: 0]

  Branch (56:24): [Folded - Ignored]

  Branch (56:24): [Folded - Ignored]

  Branch (56:24): [True: 7, False: 19]

                    return Some((Ok(map), Some(state)));

                }

                if state.data.cursor == 0 {

  Branch (59:20): [True: 0, False: 0]

  Branch (59:20): [Folded - Ignored]

  Branch (59:20): [Folded - Ignored]

  Branch (59:20): [True: 19, False: 0]

                    return None;

                }

                let data_res = state

                    .client

                    .ft_cursor_read(state.index.clone(), state.data.cursor, None)

                    .await;

                state.data = match data_res {

                    Ok(data) => data,

                    Err(err) => return Some((Err(err), None)),

                };

            }

        },

    ))

}

#[derive(Debug, Default)]

struct RedisCursorData {

    total: u64,

    cursor: u64,

    data: VecDeque<RedisMap>,

}

impl FromValue for RedisCursorData {

    fn from_value(value: RedisValue) -> Result<Self, RedisError> {

        if !value.is_array() {

  Branch (84:12): [True: 0, False: 21]

  Branch (84:12): [Folded - Ignored]

            return Err(RedisError::new(RedisErrorKind::Protocol, "Expected array"));

        }

        let mut output = Self::default();

        let value = value.into_array();

        if value.len() < 2 {

  Branch (89:12): [True: 0, False: 21]

  Branch (89:12): [Folded - Ignored]

            return Err(RedisError::new(

                RedisErrorKind::Protocol,

                "Expected at least 2 elements",

            ));

        }

        let mut value = value.into_iter();

        let data_ary = value.next().unwrap().into_array();

        if data_ary.is_empty() {

  Branch (97:12): [True: 0, False: 21]

  Branch (97:12): [Folded - Ignored]

            return Err(RedisError::new(

                RedisErrorKind::Protocol,

                "Expected at least 1 element in data array",

            ));

        }

        let Some(total) = data_ary[0].as_u64() else {

  Branch (103:13): [True: 21, False: 0]

  Branch (103:13): [Folded - Ignored]

            return Err(RedisError::new(

                RedisErrorKind::Protocol,

                "Expected integer as first element",

            ));

        };

        output.total = total;

        output.data.reserve(data_ary.len() - 1);

        for 
map_data7
 in data_ary.into_iter().skip(1) {

            output.data.push_back(map_data.into_map()
?0
);

        }

        let Some(cursor) = value.next().unwrap().as_u64() else {

  Branch (114:13): [True: 21, False: 0]

  Branch (114:13): [Folded - Ignored]

            return Err(RedisError::new(

                RedisErrorKind::Protocol,

                "Expected integer as last element",

            ));

        };

        output.cursor = cursor;

        Ok(output)

    }

}