//! Network implementation for Raft RPC
//!
//! This module provides network adapters for Raft to communicate between nodes.

use chainfire_types::NodeId;
use crate::core::{VoteRequest, VoteResponse, AppendEntriesRequest, AppendEntriesResponse};

use std::sync::Arc;
use thiserror::Error;

/// Network error type
#[derive(Error, Debug)]
pub enum RaftNetworkError {
    #[error("Connection failed to node {node_id}: {reason}")]
    ConnectionFailed { node_id: NodeId, reason: String },

    #[error("RPC failed: {0}")]
    RpcFailed(String),

    #[error("Timeout")]
    Timeout,

    #[error("Node {0} not found")]
    NodeNotFound(NodeId),
}

/// Trait for sending Raft RPCs
#[async_trait::async_trait]
pub trait RaftRpcClient: Send + Sync + 'static {
    async fn vote(
        &self,
        target: NodeId,
        req: VoteRequest,
    ) -> Result<VoteResponse, RaftNetworkError>;

    async fn append_entries(
        &self,
        target: NodeId,
        req: AppendEntriesRequest,
    ) -> Result<AppendEntriesResponse, RaftNetworkError>;
}

/// In-memory RPC client for testing
pub mod test_client {
    use super::*;
    use std::collections::HashMap;
    use tokio::sync::mpsc;

    /// A simple in-memory RPC client for testing custom Raft
    #[derive(Clone)]
    pub struct InMemoryRpcClient {
        /// Channel senders to each node
        channels: Arc<tokio::sync::RwLock<HashMap<NodeId, mpsc::UnboundedSender<RpcMessage>>>>,
    }

    pub enum RpcMessage {
        Vote(
            VoteRequest,
            tokio::sync::oneshot::Sender<VoteResponse>,
        ),
        AppendEntries(
            AppendEntriesRequest,
            tokio::sync::oneshot::Sender<AppendEntriesResponse>,
        ),
    }

    impl Default for InMemoryRpcClient {
        fn default() -> Self {
            Self::new()
        }
    }

    impl InMemoryRpcClient {
        pub fn new() -> Self {
            Self {
                channels: Arc::new(tokio::sync::RwLock::new(HashMap::new())),
            }
        }

        pub async fn register(&self, id: NodeId, tx: mpsc::UnboundedSender<RpcMessage>) {
            self.channels.write().await.insert(id, tx);
        }
    }

    #[async_trait::async_trait]
    impl RaftRpcClient for InMemoryRpcClient {
        async fn vote(
            &self,
            target: NodeId,
            req: VoteRequest,
        ) -> Result<VoteResponse, RaftNetworkError> {
            let channels = self.channels.read().await;
            let tx = channels
                .get(&target)
                .ok_or(RaftNetworkError::NodeNotFound(target))?;

            let (resp_tx, resp_rx) = tokio::sync::oneshot::channel();
            tx.send(RpcMessage::Vote(req, resp_tx))
                .map_err(|_| RaftNetworkError::RpcFailed("Channel closed".into()))?;

            resp_rx
                .await
                .map_err(|_| RaftNetworkError::RpcFailed("Response channel closed".into()))
        }

        async fn append_entries(
            &self,
            target: NodeId,
            req: AppendEntriesRequest,
        ) -> Result<AppendEntriesResponse, RaftNetworkError> {
            let channels = self.channels.read().await;
            let tx = channels
                .get(&target)
                .ok_or_else(|| {
                    eprintln!("[RPC] NodeNotFound: target={}, registered={:?}",
                        target, channels.keys().collect::<Vec<_>>());
                    RaftNetworkError::NodeNotFound(target)
                })?;

            let (resp_tx, resp_rx) = tokio::sync::oneshot::channel();
            let send_result = tx.send(RpcMessage::AppendEntries(req.clone(), resp_tx));

            if let Err(e) = send_result {
                eprintln!("[RPC] Send failed to node {}: channel closed", target);
                return Err(RaftNetworkError::RpcFailed("Channel closed".into()));
            }

            resp_rx
                .await
                .map_err(|_| RaftNetworkError::RpcFailed("Response channel closed".into()))
        }
    }
}
