//! Chainfire client implementation

use crate::error::{ClientError, Result};
use crate::watch::WatchHandle;
use chainfire_proto::proto::{
    cluster_client::ClusterClient, compare, kv_client::KvClient, request_op, response_op,
    watch_client::WatchClient, Compare, DeleteRangeRequest, MemberAddRequest, PutRequest,
    RangeRequest, RequestOp, StatusRequest, TxnRequest,
};
use std::time::Duration;
use tonic::transport::Channel;
use tonic::Code;
use tracing::{debug, warn};

/// Chainfire client
pub struct Client {
    /// Configured client endpoints
    endpoints: Vec<String>,
    /// Preferred endpoint index
    current_endpoint: usize,
    /// gRPC channel
    channel: Channel,
    /// KV client
    kv: KvClient<Channel>,
    /// Cluster client
    cluster: ClusterClient<Channel>,
}

impl Client {
    /// Connect to a Chainfire server
    pub async fn connect(addr: impl AsRef<str>) -> Result<Self> {
        let endpoints = parse_endpoints(addr.as_ref())?;
        let mut last_error = None;

        for (index, endpoint) in endpoints.iter().enumerate() {
            match connect_endpoint(endpoint).await {
                Ok((channel, kv, cluster)) => {
                    debug!(endpoint = %endpoint, "Connected to Chainfire");
                    let mut client = Self {
                        endpoints: endpoints.clone(),
                        current_endpoint: index,
                        channel,
                        kv,
                        cluster,
                    };
                    client.promote_leader_endpoint().await?;
                    return Ok(client);
                }
                Err(error) => {
                    warn!(endpoint = %endpoint, error = %error, "Chainfire endpoint connect failed");
                    last_error = Some(error);
                }
            }
        }

        Err(last_error.unwrap_or_else(|| {
            ClientError::Connection("no Chainfire endpoints configured".to_string())
        }))
    }

    async fn with_kv_retry<T, F, Fut>(&mut self, mut op: F) -> Result<T>
    where
        F: FnMut(KvClient<Channel>) -> Fut,
        Fut: std::future::Future<Output = std::result::Result<T, tonic::Status>>,
    {
        let max_attempts = self.endpoints.len().max(1) * 3;
        let mut last_status = None;
        for attempt in 0..max_attempts {
            let client = self.kv.clone();
            match op(client).await {
                Ok(value) => return Ok(value),
                Err(status) if attempt + 1 < max_attempts && is_retryable_status(&status) => {
                    warn!(
                        endpoint = %self.endpoints[self.current_endpoint],
                        code = ?status.code(),
                        message = %status.message(),
                        attempt = attempt + 1,
                        max_attempts,
                        "retrying Chainfire KV RPC on alternate endpoint"
                    );
                    last_status = Some(status);
                    self.recover_after_status(last_status.as_ref().unwrap())
                        .await?;
                    tokio::time::sleep(retry_delay(attempt)).await;
                }
                Err(status) => return Err(status.into()),
            }
        }

        Err(last_status
            .unwrap_or_else(|| tonic::Status::unavailable("Chainfire KV retry exhausted"))
            .into())
    }

    async fn with_cluster_retry<T, F, Fut>(&mut self, mut op: F) -> Result<T>
    where
        F: FnMut(ClusterClient<Channel>) -> Fut,
        Fut: std::future::Future<Output = std::result::Result<T, tonic::Status>>,
    {
        let max_attempts = self.endpoints.len().max(1) * 3;
        let mut last_status = None;
        for attempt in 0..max_attempts {
            let client = self.cluster.clone();
            match op(client).await {
                Ok(value) => return Ok(value),
                Err(status) if attempt + 1 < max_attempts && is_retryable_status(&status) => {
                    warn!(
                        endpoint = %self.endpoints[self.current_endpoint],
                        code = ?status.code(),
                        message = %status.message(),
                        attempt = attempt + 1,
                        max_attempts,
                        "retrying Chainfire cluster RPC on alternate endpoint"
                    );
                    last_status = Some(status);
                    self.recover_after_status(last_status.as_ref().unwrap())
                        .await?;
                    tokio::time::sleep(retry_delay(attempt)).await;
                }
                Err(status) => return Err(status.into()),
            }
        }

        Err(last_status
            .unwrap_or_else(|| tonic::Status::unavailable("Chainfire cluster retry exhausted"))
            .into())
    }

    async fn recover_after_status(&mut self, status: &tonic::Status) -> Result<()> {
        if let Some(leader_idx) = self.discover_leader_endpoint().await? {
            if leader_idx != self.current_endpoint {
                return self.reconnect_to_index(leader_idx).await;
            }
        }

        if self.endpoints.len() > 1 {
            let next = (self.current_endpoint + 1) % self.endpoints.len();
            if next != self.current_endpoint {
                return self.reconnect_to_index(next).await;
            }
        }

        Err(ClientError::Rpc(status.clone()))
    }

    async fn reconnect_to_index(&mut self, index: usize) -> Result<()> {
        let endpoint = self
            .endpoints
            .get(index)
            .ok_or_else(|| {
                ClientError::Connection(format!("invalid Chainfire endpoint index {index}"))
            })?
            .clone();
        let (channel, kv, cluster) = connect_endpoint(&endpoint).await?;
        self.current_endpoint = index;
        self.channel = channel;
        self.kv = kv;
        self.cluster = cluster;
        Ok(())
    }

    async fn promote_leader_endpoint(&mut self) -> Result<()> {
        if let Some(index) = self.discover_leader_endpoint().await? {
            if index != self.current_endpoint {
                self.reconnect_to_index(index).await?;
            }
        }
        Ok(())
    }

    async fn discover_leader_endpoint(&self) -> Result<Option<usize>> {
        for (index, endpoint) in self.endpoints.iter().enumerate() {
            let mut cluster = match ClusterClient::connect(endpoint.clone()).await {
                Ok(client) => client,
                Err(error) => {
                    warn!(endpoint = %endpoint, error = %error, "failed to connect while probing Chainfire leader");
                    continue;
                }
            };

            match cluster.status(StatusRequest {}).await {
                Ok(response) => {
                    let status = response.into_inner();
                    let member_id = status
                        .header
                        .as_ref()
                        .map(|header| header.member_id)
                        .unwrap_or(0);
                    if status.leader != 0 && status.leader == member_id {
                        return Ok(Some(index));
                    }
                }
                Err(status) => {
                    warn!(
                        endpoint = %endpoint,
                        code = ?status.code(),
                        message = %status.message(),
                        "failed to query Chainfire leader status"
                    );
                }
            }
        }

        Ok(None)
    }

    /// Put a key-value pair
    pub async fn put(&mut self, key: impl AsRef<[u8]>, value: impl AsRef<[u8]>) -> Result<u64> {
        let key = key.as_ref().to_vec();
        let value = value.as_ref().to_vec();
        let resp = self
            .with_kv_retry(|mut kv| {
                let key = key.clone();
                let value = value.clone();
                async move {
                    kv.put(PutRequest {
                        key,
                        value,
                        lease: 0,
                        prev_kv: false,
                    })
                    .await
                    .map(|resp| resp.into_inner())
                }
            })
            .await?;

        Ok(resp.header.map(|h| h.revision as u64).unwrap_or(0))
    }

    /// Put a key-value pair with string values
    pub async fn put_str(&mut self, key: &str, value: &str) -> Result<u64> {
        self.put(key.as_bytes(), value.as_bytes()).await
    }

    /// Get a value by key
    pub async fn get(&mut self, key: impl AsRef<[u8]>) -> Result<Option<Vec<u8>>> {
        Ok(self.get_with_revision(key).await?.map(|(value, _)| value))
    }

    /// Get a value by key along with its current revision
    pub async fn get_with_revision(
        &mut self,
        key: impl AsRef<[u8]>,
    ) -> Result<Option<(Vec<u8>, u64)>> {
        let key = key.as_ref().to_vec();
        let resp = self
            .with_kv_retry(|mut kv| {
                let key = key.clone();
                async move {
                    kv.range(RangeRequest {
                        key,
                        range_end: vec![],
                        limit: 1,
                        revision: 0,
                        keys_only: false,
                        count_only: false,
                        serializable: false,
                    })
                    .await
                    .map(|resp| resp.into_inner())
                }
            })
            .await?;

        Ok(resp
            .kvs
            .into_iter()
            .next()
            .map(|kv| (kv.value, kv.mod_revision as u64)))
    }

    /// Put a key-value pair only if the key's mod_revision matches.
    pub async fn put_if_revision(
        &mut self,
        key: impl AsRef<[u8]>,
        value: impl AsRef<[u8]>,
        expected_mod_revision: u64,
    ) -> Result<()> {
        let key_bytes = key.as_ref().to_vec();
        let key_display = String::from_utf8_lossy(&key_bytes).to_string();
        let compare = Compare {
            result: compare::CompareResult::Equal as i32,
            target: compare::CompareTarget::Mod as i32,
            key: key_bytes.clone(),
            target_union: Some(compare::TargetUnion::ModRevision(
                expected_mod_revision as i64,
            )),
        };

        let put_op = RequestOp {
            request: Some(request_op::Request::RequestPut(PutRequest {
                key: key_bytes.clone(),
                value: value.as_ref().to_vec(),
                lease: 0,
                prev_kv: false,
            })),
        };

        let read_on_fail = RequestOp {
            request: Some(request_op::Request::RequestRange(RangeRequest {
                key: key_bytes.clone(),
                range_end: vec![],
                limit: 1,
                revision: 0,
                keys_only: false,
                count_only: false,
                serializable: false,
            })),
        };

        let resp = self
            .with_kv_retry(|mut kv| {
                let compare = compare.clone();
                let put_op = put_op.clone();
                let read_on_fail = read_on_fail.clone();
                async move {
                    kv.txn(TxnRequest {
                        compare: vec![compare],
                        success: vec![put_op],
                        failure: vec![read_on_fail],
                    })
                    .await
                    .map(|resp| resp.into_inner())
                }
            })
            .await?;

        if resp.succeeded {
            return Ok(());
        }

        let current_revision = resp
            .responses
            .into_iter()
            .filter_map(|op| match op.response {
                Some(response_op::Response::ResponseRange(range)) => range
                    .kvs
                    .into_iter()
                    .next()
                    .map(|kv| kv.mod_revision as u64),
                _ => None,
            })
            .next()
            .unwrap_or(0);

        Err(ClientError::Conflict(format!(
            "mod_revision mismatch for key {key_display}: expected {expected_mod_revision}, current {current_revision}"
        )))
    }

    /// Get a value as string
    pub async fn get_str(&mut self, key: &str) -> Result<Option<String>> {
        let value = self.get(key.as_bytes()).await?;
        Ok(value.map(|v| String::from_utf8_lossy(&v).to_string()))
    }

    /// Delete a key
    pub async fn delete(&mut self, key: impl AsRef<[u8]>) -> Result<bool> {
        let key = key.as_ref().to_vec();
        let resp = self
            .with_kv_retry(|mut kv| {
                let key = key.clone();
                async move {
                    kv.delete(DeleteRangeRequest {
                        key,
                        range_end: vec![],
                        prev_kv: false,
                    })
                    .await
                    .map(|resp| resp.into_inner())
                }
            })
            .await?;

        Ok(resp.deleted > 0)
    }

    /// Get all keys with a prefix
    pub async fn get_prefix(
        &mut self,
        prefix: impl AsRef<[u8]>,
    ) -> Result<Vec<(Vec<u8>, Vec<u8>)>> {
        let prefix = prefix.as_ref();
        let range_end = prefix_end(prefix);

        let resp = self
            .with_kv_retry(|mut kv| {
                let key = prefix.to_vec();
                let range_end = range_end.clone();
                async move {
                    kv.range(RangeRequest {
                        key,
                        range_end,
                        limit: 0,
                        revision: 0,
                        keys_only: false,
                        count_only: false,
                        serializable: false,
                    })
                    .await
                    .map(|resp| resp.into_inner())
                }
            })
            .await?;

        Ok(resp.kvs.into_iter().map(|kv| (kv.key, kv.value)).collect())
    }

    /// Scan a prefix returning keys, values, and revisions
    pub async fn scan_prefix(
        &mut self,
        prefix: impl AsRef<[u8]>,
        limit: i64,
    ) -> Result<(Vec<(Vec<u8>, Vec<u8>, u64)>, Option<Vec<u8>>)> {
        let prefix = prefix.as_ref();
        let range_end = prefix_end(prefix);

        let resp = self
            .with_kv_retry(|mut kv| {
                let key = prefix.to_vec();
                let range_end = range_end.clone();
                async move {
                    kv.range(RangeRequest {
                        key,
                        range_end,
                        limit,
                        revision: 0,
                        keys_only: false,
                        count_only: false,
                        serializable: false,
                    })
                    .await
                    .map(|resp| resp.into_inner())
                }
            })
            .await?;

        let more = resp.more;
        let kvs: Vec<(Vec<u8>, Vec<u8>, u64)> = resp
            .kvs
            .into_iter()
            .map(|kv| (kv.key, kv.value, kv.mod_revision as u64))
            .collect();
        let next_key = if more {
            kvs.last().map(|(k, _, _)| {
                let mut nk = k.clone();
                nk.push(0);
                nk
            })
        } else {
            None
        };

        Ok((kvs, next_key))
    }

    /// Scan an arbitrary range [start, end)
    pub async fn scan_range(
        &mut self,
        start: impl AsRef<[u8]>,
        end: impl AsRef<[u8]>,
        limit: i64,
    ) -> Result<(Vec<(Vec<u8>, Vec<u8>, u64)>, Option<Vec<u8>>)> {
        let resp = self
            .with_kv_retry(|mut kv| {
                let key = start.as_ref().to_vec();
                let range_end = end.as_ref().to_vec();
                async move {
                    kv.range(RangeRequest {
                        key,
                        range_end,
                        limit,
                        revision: 0,
                        keys_only: false,
                        count_only: false,
                        serializable: false,
                    })
                    .await
                    .map(|resp| resp.into_inner())
                }
            })
            .await?;

        let more = resp.more;
        let kvs: Vec<(Vec<u8>, Vec<u8>, u64)> = resp
            .kvs
            .into_iter()
            .map(|kv| (kv.key, kv.value, kv.mod_revision as u64))
            .collect();
        let next_key = if more {
            kvs.last().map(|(k, _, _)| {
                let mut nk = k.clone();
                nk.push(0);
                nk
            })
        } else {
            None
        };

        Ok((kvs, next_key))
    }

    /// Compare-and-swap based on key version
    pub async fn compare_and_swap(
        &mut self,
        key: impl AsRef<[u8]>,
        expected_version: u64,
        value: impl AsRef<[u8]>,
    ) -> Result<CasOutcome> {
        let key_bytes = key.as_ref().to_vec();
        let put_op = RequestOp {
            request: Some(request_op::Request::RequestPut(PutRequest {
                key: key_bytes.clone(),
                value: value.as_ref().to_vec(),
                lease: 0,
                prev_kv: false,
            })),
        };

        // Fetch current value on failure to surface the actual version
        let read_on_fail = RequestOp {
            request: Some(request_op::Request::RequestRange(RangeRequest {
                key: key_bytes.clone(),
                range_end: vec![],
                limit: 1,
                revision: 0,
                keys_only: false,
                count_only: false,
                serializable: true, // within txn, use serializable read
            })),
        };

        let compare = Compare {
            result: compare::CompareResult::Equal as i32,
            target: compare::CompareTarget::Version as i32,
            key: key_bytes.clone(),
            target_union: Some(compare::TargetUnion::Version(expected_version as i64)),
        };

        let resp = self
            .with_kv_retry(|mut kv| {
                let compare = compare.clone();
                let put_op = put_op.clone();
                let read_on_fail = read_on_fail.clone();
                async move {
                    kv.txn(TxnRequest {
                        compare: vec![compare],
                        success: vec![put_op],
                        failure: vec![read_on_fail],
                    })
                    .await
                    .map(|resp| resp.into_inner())
                }
            })
            .await?;

        if resp.succeeded {
            let new_version = resp.header.as_ref().map(|h| h.revision as u64).unwrap_or(0);
            return Ok(CasOutcome {
                success: true,
                current_version: new_version,
                new_version,
            });
        }

        // On failure try to extract the current version from the range response
        let current_version = resp
            .responses
            .into_iter()
            .filter_map(|op| match op.response {
                Some(response_op::Response::ResponseRange(r)) => {
                    r.kvs.into_iter().next().map(|kv| kv.mod_revision as u64)
                }
                _ => None,
            })
            .next()
            .unwrap_or(0);

        Ok(CasOutcome {
            success: false,
            current_version,
            new_version: 0,
        })
    }

    /// Watch a key or prefix for changes
    pub async fn watch(&mut self, key: impl AsRef<[u8]>) -> Result<WatchHandle> {
        let key = key.as_ref().to_vec();
        let watch_client = WatchClient::new(self.channel.clone());
        WatchHandle::new(watch_client, key, None).await
    }

    /// Watch all keys with a prefix
    pub async fn watch_prefix(&mut self, prefix: impl AsRef<[u8]>) -> Result<WatchHandle> {
        let prefix = prefix.as_ref().to_vec();
        let range_end = prefix_end(&prefix);
        let watch_client = WatchClient::new(self.channel.clone());
        WatchHandle::new(watch_client, prefix, Some(range_end)).await
    }

    /// Get cluster status
    pub async fn status(&mut self) -> Result<ClusterStatus> {
        let resp = self
            .with_cluster_retry(|mut cluster| async move {
                cluster
                    .status(StatusRequest {})
                    .await
                    .map(|resp| resp.into_inner())
            })
            .await?;

        Ok(ClusterStatus {
            version: resp.version,
            leader: resp.leader,
            raft_term: resp.raft_term,
        })
    }

    /// Add a member to the cluster
    ///
    /// # Arguments
    /// * `peer_url` - The Raft address of the new member (e.g., "127.0.0.1:2380")
    /// * `is_learner` - Whether to add as learner (true) or voter (false)
    ///
    /// # Returns
    /// The node ID of the added member
    pub async fn member_add(
        &mut self,
        node_id: u64,
        peer_url: impl AsRef<str>,
        is_learner: bool,
    ) -> Result<u64> {
        let peer_url = peer_url.as_ref().to_string();
        let resp = self
            .with_cluster_retry(|mut cluster| {
                let peer_url = peer_url.clone();
                async move {
                    cluster
                        .member_add(MemberAddRequest {
                            node_id,
                            peer_urls: vec![peer_url],
                            is_learner,
                        })
                        .await
                        .map(|resp| resp.into_inner())
                }
            })
            .await?;

        // Extract the member ID from the response
        let member_id = resp
            .member
            .map(|m| m.id)
            .ok_or_else(|| ClientError::Internal("No member in response".to_string()))?;

        debug!(
            member_id = member_id,
            peer_url = peer_url.as_str(),
            is_learner = is_learner,
            "Added member to cluster"
        );

        Ok(member_id)
    }
}

/// Cluster status
#[derive(Debug, Clone)]
pub struct ClusterStatus {
    /// Server version
    pub version: String,
    /// Current leader ID
    pub leader: u64,
    /// Current Raft term
    pub raft_term: u64,
}

/// CAS outcome returned by compare_and_swap
#[derive(Debug, Clone)]
pub struct CasOutcome {
    /// Whether CAS succeeded
    pub success: bool,
    /// Observed/current version
    pub current_version: u64,
    /// New version when succeeded
    pub new_version: u64,
}

fn parse_endpoints(input: &str) -> Result<Vec<String>> {
    let endpoints: Vec<String> = input
        .split(',')
        .map(str::trim)
        .filter(|value| !value.is_empty())
        .map(normalize_endpoint)
        .collect();

    if endpoints.is_empty() {
        return Err(ClientError::Connection(
            "no Chainfire endpoints configured".to_string(),
        ));
    }

    Ok(endpoints)
}

fn normalize_endpoint(endpoint: &str) -> String {
    if endpoint.contains("://") {
        endpoint.to_string()
    } else {
        format!("http://{endpoint}")
    }
}

async fn connect_endpoint(
    endpoint: &str,
) -> Result<(Channel, KvClient<Channel>, ClusterClient<Channel>)> {
    let channel = Channel::from_shared(endpoint.to_string())
        .map_err(|e| ClientError::Connection(e.to_string()))?
        .connect()
        .await?;

    let kv = KvClient::new(channel.clone());
    let cluster = ClusterClient::new(channel.clone());
    Ok((channel, kv, cluster))
}

fn retry_delay(attempt: usize) -> Duration {
    let multiplier = 1u64 << attempt.min(3);
    Duration::from_millis((200 * multiplier).min(1_000))
}

fn is_retryable_status(status: &tonic::Status) -> bool {
    matches!(
        status.code(),
        Code::Unavailable
            | Code::DeadlineExceeded
            | Code::Internal
            | Code::Aborted
            | Code::FailedPrecondition
    ) || retryable_message(status.message())
}

fn retryable_message(message: &str) -> bool {
    let lowercase = message.to_ascii_lowercase();
    lowercase.contains("not leader")
        || lowercase.contains("leader_id")
        || lowercase.contains("transport error")
        || lowercase.contains("connection was not ready")
        || lowercase.contains("deadline has elapsed")
        || lowercase.contains("broken pipe")
        || lowercase.contains("connection reset")
        || lowercase.contains("connection refused")
}

/// Calculate prefix end for range queries
fn prefix_end(prefix: &[u8]) -> Vec<u8> {
    let mut end = prefix.to_vec();
    for i in (0..end.len()).rev() {
        if end[i] < 0xff {
            end[i] += 1;
            end.truncate(i + 1);
            return end;
        }
    }
    vec![]
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_prefix_end() {
        assert_eq!(prefix_end(b"abc"), b"abd");
        assert_eq!(prefix_end(b"/nodes/"), b"/nodes0");
    }

    #[test]
    fn normalize_endpoint_adds_http_scheme() {
        assert_eq!(
            normalize_endpoint("127.0.0.1:2379"),
            "http://127.0.0.1:2379"
        );
        assert_eq!(
            normalize_endpoint("http://127.0.0.1:2379"),
            "http://127.0.0.1:2379"
        );
    }

    #[test]
    fn parse_endpoints_accepts_comma_separated_values() {
        let endpoints = parse_endpoints("127.0.0.1:2379, http://127.0.0.2:2379").unwrap();
        assert_eq!(
            endpoints,
            vec![
                "http://127.0.0.1:2379".to_string(),
                "http://127.0.0.2:2379".to_string()
            ]
        );
    }

    #[test]
    fn retryable_message_covers_not_leader_and_transport() {
        assert!(retryable_message("NotLeader { leader_id: Some(1) }"));
        assert!(retryable_message("transport error"));
        assert!(retryable_message("connection was not ready"));
        assert!(!retryable_message("permission denied"));
    }
}
