rostrum 14.0.1

use anyhow::Context;
use jsonrpc::simple_http::SimpleHttpTransport;
use serde_json::value::{to_raw_value, RawValue};
use std::net::SocketAddr;
use std::ops::Deref;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use tokio::sync::Mutex;

use jsonrpc::{Client, Request};
use serde_json::Value;

use crate::daemon::CookieGetter;
use crate::errors::{rpc_invalid_params, rpc_other, ConnectionError};

pub struct RPCConnection {
    client: Mutex<Option<Arc<jsonrpc::Client>>>,
    is_broken: AtomicBool,
}

impl RPCConnection {
    pub(crate) fn new(
        addr: SocketAddr,
        cookie_getter: Arc<dyn CookieGetter>,
    ) -> Result<Self, ConnectionError> {
        let cookie = match cookie_getter.get() {
            Ok(c) => c,
            Err(e) => {
                return Err(ConnectionError {
                    msg: format!("Authentication error: {e}"),
                })
            }
        };

        let t = match SimpleHttpTransport::builder().url(&addr.to_string()) {
            Ok(t) => t,
            Err(e) => {
                return Err(ConnectionError {
                    msg: format!("Invalid RPC address: {e}"),
                })
            }
        };
        let t = t.cookie_auth(String::from_utf8(cookie).unwrap()).build();

        Ok(Self {
            client: Mutex::new(Some(Arc::new(Client::with_transport(t)))),
            is_broken: AtomicBool::new(false),
        })
    }

    async fn check_error(&self, response_error: jsonrpc::Error, method: &str) -> anyhow::Error {
        match response_error {
            jsonrpc::Error::Rpc(rpc_error) => {
                match rpc_error.code {
                    // RPC_IN_WARMUP -> retry by later reconnection.
                    // We don't use "self.connection_error" to flag that the connection is broken,
                    // because the RPC server is simply not ready.
                    -28 => anyhow::Error::from(ConnectionError {
                        msg: rpc_error.message,
                    }),
                    // RPC_INVALID_ADDRESS_OR_KEY
                    -5 => rpc_invalid_params(rpc_error.message),
                    _ => rpc_other(format!(
                        "Call '{}' to full node failed: {}",
                        method, rpc_error.message
                    )),
                }
            }
            _ => anyhow::Error::from(self.connection_error(response_error.to_string()).await),
        }
    }

    pub async fn request(&self, method: &str, params: &Option<Value>) -> anyhow::Result<Value> {
        let client = match self.client.lock().await.as_ref() {
            Some(c) => Arc::clone(c),
            None => {
                return Err(anyhow::Error::msg(ConnectionError {
                    msg: "no client".to_string(),
                }))
            }
        };

        let method_cpy = method.to_string();
        let params_cpy = match params {
            Some(v) => Some(to_raw_value(&v)?),
            None => None,
        };

        let response = tokio::task::spawn_blocking(move || {
            let unboxed = params_cpy.as_deref();
            let request = client.build_request(&method_cpy, unboxed);
            client.send_request(request)
        })
        .await?;

        let response = match response {
            Ok(r) => Ok(r),
            Err(e) => Err(anyhow::Error::msg(
                self.connection_error(e.to_string()).await,
            )),
        }?;

        match response.result() {
            Ok(r) => Ok(r),
            Err(e) => Err(self.check_error(e, method).await),
        }
    }

    async fn connection_error(&self, what: String) -> ConnectionError {
        trace!("Flagging a connection as broken due to: {what}");

        // Lift the "is broken" flag used by connection pool to determine if connection is good.
        self.is_broken.store(true, Ordering::Relaxed);

        // drop our ref to this connection since it has issues
        self.client.lock().await.take();

        ConnectionError { msg: what }
    }

    /**
     * Flag will be true of ConnectionError was cast at some point.
     */
    pub fn is_broken(&self) -> bool {
        self.is_broken.load(Ordering::Relaxed)
    }

    /**
     * Do multiple request with different parameters for same method.
     */
    pub async fn multi_request(
        &self,
        method: &str,
        params_list: &[Option<Value>],
    ) -> anyhow::Result<Vec<Value>> {
        let method_cpy = method.to_string();
        let mut params_list_cpy: Vec<Option<Box<RawValue>>> = Vec::with_capacity(params_list.len());

        for p in params_list {
            let param = match p {
                Some(p) => Some(to_raw_value(p)?),
                None => None,
            };
            params_list_cpy.push(param)
        }

        let client = match self.client.lock().await.as_ref() {
            Some(c) => Arc::clone(c),
            None => {
                return Err(anyhow::Error::msg(ConnectionError {
                    msg: "no client".to_string(),
                }))
            }
        };
        let (responses, reqs_len) = tokio::task::spawn_blocking(move || {
            let unboxed: Vec<Option<&RawValue>> = params_list_cpy
                .iter()
                .map(|v| {
                    let v = v.as_ref().map(|v| v.deref());
                    v
                })
                .collect();

            let reqs: Vec<Request> = unboxed
                .into_iter()
                .map(|params| client.build_request(&method_cpy, params))
                .collect();
            let reqs_len = reqs.len();
            let responses = client.send_batch(&reqs);
            (responses, reqs_len)
        })
        .await?;

        let responses = match responses {
            Ok(r) => r,
            Err(e) => {
                return Err(anyhow::Error::msg(
                    self.connection_error(e.to_string()).await,
                ))
            }
        };

        let mut result: Vec<Value> = Vec::with_capacity(reqs_len);
        for reply in responses {
            let response = reply.context("Received empty response to {method}")?;
            match response.result() {
                Ok(r) => result.push(r),
                Err(e) => return Err(self.check_error(e, method).await),
            }
        }
        Ok(result)
    }
}