//! RaftNetwork implementation for FlareDB
//!
//! This module implements OpenRaft's RaftNetwork and RaftNetworkFactory traits
//! for inter-node communication using gRPC.

use crate::types::{FlareNode, FlareNodeId, FlareTypeConfig};
use openraft::error::{Fatal, NetworkError, RPCError, RaftError, ReplicationClosed, StreamingError};
use openraft::network::{RPCOption, RaftNetwork, RaftNetworkFactory};
use openraft::raft::{
    AppendEntriesRequest, AppendEntriesResponse, InstallSnapshotRequest, InstallSnapshotResponse,
    SnapshotResponse, VoteRequest, VoteResponse,
};
use openraft::{OptionalSend, Snapshot, Vote};
use flaredb_proto::raft_server::raft_service_client::RaftServiceClient;
use std::collections::HashMap;
use std::future::Future;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::RwLock;
use tonic::transport::{Channel, Endpoint};
use tracing::{debug, error};

const RAFT_GRPC_MESSAGE_SIZE: usize = 64 * 1024 * 1024;
const RAFT_CONNECT_TIMEOUT: Duration = Duration::from_secs(2);
const RAFT_GRPC_INITIAL_STREAM_WINDOW: u32 = 16 * 1024 * 1024;
const RAFT_GRPC_INITIAL_CONNECTION_WINDOW: u32 = 64 * 1024 * 1024;
const RAFT_GRPC_KEEPALIVE_INTERVAL: Duration = Duration::from_secs(5);
const RAFT_GRPC_KEEPALIVE_TIMEOUT: Duration = Duration::from_secs(2);

/// Type aliases for cleaner code
type FlareRPCError = RPCError<FlareNodeId, FlareNode, RaftError<FlareNodeId>>;
type FlareSnapshotError = StreamingError<FlareTypeConfig, Fatal<FlareNodeId>>;
type FlareInstallSnapshotError =
    RPCError<FlareNodeId, FlareNode, RaftError<FlareNodeId, openraft::error::InstallSnapshotError>>;

/// Network factory for creating per-node connections
pub struct FlareNetworkFactory {
    local_id: FlareNodeId,
    region_id: u64,
    clients: Arc<RwLock<HashMap<FlareNodeId, RaftServiceClient<Channel>>>>,
    node_addrs: Arc<RwLock<HashMap<FlareNodeId, String>>>,
}

impl FlareNetworkFactory {
    pub fn new(local_id: FlareNodeId, region_id: u64) -> Self {
        Self {
            local_id,
            region_id,
            clients: Arc::new(RwLock::new(HashMap::new())),
            node_addrs: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    /// Add a peer's address to the registry
    pub async fn add_peer(&self, peer_id: FlareNodeId, addr: String) {
        let mut addrs = self.node_addrs.write().await;
        addrs.insert(peer_id, addr);
    }

    /// Remove a peer from the registry
    pub async fn remove_peer(&self, peer_id: FlareNodeId) {
        let mut addrs = self.node_addrs.write().await;
        addrs.remove(&peer_id);

        let mut clients = self.clients.write().await;
        clients.remove(&peer_id);
    }

    /// Get all known peer addresses
    pub async fn get_peers(&self) -> HashMap<FlareNodeId, String> {
        self.node_addrs.read().await.clone()
    }
}

impl Clone for FlareNetworkFactory {
    fn clone(&self) -> Self {
        Self {
            local_id: self.local_id,
            region_id: self.region_id,
            clients: self.clients.clone(),
            node_addrs: self.node_addrs.clone(),
        }
    }
}

impl RaftNetworkFactory<FlareTypeConfig> for FlareNetworkFactory {
    type Network = FlareNetwork;

    async fn new_client(&mut self, target: FlareNodeId, node: &FlareNode) -> Self::Network {
        // Register the node address
        if !node.addr.is_empty() {
            let mut addrs = self.node_addrs.write().await;
            addrs.insert(target, node.addr.clone());
        }

        FlareNetwork {
            target,
            region_id: self.region_id,
            clients: self.clients.clone(),
            node_addrs: self.node_addrs.clone(),
        }
    }
}

/// Network client for a specific remote node
pub struct FlareNetwork {
    target: FlareNodeId,
    region_id: u64,
    clients: Arc<RwLock<HashMap<FlareNodeId, RaftServiceClient<Channel>>>>,
    node_addrs: Arc<RwLock<HashMap<FlareNodeId, String>>>,
}

impl FlareNetwork {
    fn endpoint_builder(addr: &str) -> Result<Endpoint, String> {
        Endpoint::from_shared(format!("http://{}", addr))
            .map_err(|e| e.to_string())
            .map(|endpoint| {
                endpoint
                    .connect_timeout(RAFT_CONNECT_TIMEOUT)
                    .tcp_nodelay(true)
                    .http2_keep_alive_interval(RAFT_GRPC_KEEPALIVE_INTERVAL)
                    .keep_alive_timeout(RAFT_GRPC_KEEPALIVE_TIMEOUT)
                    .initial_stream_window_size(RAFT_GRPC_INITIAL_STREAM_WINDOW)
                    .initial_connection_window_size(RAFT_GRPC_INITIAL_CONNECTION_WINDOW)
            })
    }

    /// Get or create a gRPC client for the target node
    async fn get_client(&self) -> Result<RaftServiceClient<Channel>, String> {
        // Check cached client
        {
            let clients = self.clients.read().await;
            if let Some(client) = clients.get(&self.target) {
                return Ok(client.clone());
            }
        }

        // Get target address
        let addr = {
            let addrs = self.node_addrs.read().await;
            addrs.get(&self.target).cloned()
        };

        let addr = addr.ok_or_else(|| format!("address not known for node {}", self.target))?;

        // Create new connection
        let channel = Self::endpoint_builder(&addr)?
            .connect()
            .await
            .map_err(|e| e.to_string())?;

        let client = RaftServiceClient::new(channel)
            .max_decoding_message_size(RAFT_GRPC_MESSAGE_SIZE)
            .max_encoding_message_size(RAFT_GRPC_MESSAGE_SIZE);

        // Cache the client
        let mut clients = self.clients.write().await;
        clients.insert(self.target, client.clone());

        Ok(client)
    }

    async fn evict_client(&self) {
        let mut clients = self.clients.write().await;
        clients.remove(&self.target);
    }

    fn network_error(msg: impl ToString) -> NetworkError {
        NetworkError::new(&std::io::Error::new(
            std::io::ErrorKind::ConnectionRefused,
            msg.to_string(),
        ))
    }
}

impl RaftNetwork<FlareTypeConfig> for FlareNetwork {
    async fn append_entries(
        &mut self,
        rpc: AppendEntriesRequest<FlareTypeConfig>,
        _option: RPCOption,
    ) -> Result<AppendEntriesResponse<FlareNodeId>, FlareRPCError> {
        debug!(
            target = self.target,
            region = self.region_id,
            entries = rpc.entries.len(),
            "sending append_entries"
        );

        let mut client = self
            .get_client()
            .await
            .map_err(|e| RPCError::Network(Self::network_error(e)))?;

        // Serialize the request
        let serialized = serde_json::to_vec(&rpc)
            .map_err(|e| RPCError::Network(Self::network_error(e.to_string())))?;

        // Send via gRPC using OpenRaft message type
        let grpc_req = flaredb_proto::raft_server::OpenRaftAppendEntriesRequest {
            region_id: self.region_id,
            data: serialized,
        };

        let response = client.append_entries_v2(grpc_req).await.map_err(|e| {
            error!(error = %e, "append_entries RPC failed");
            let target = self.target;
            let clients = Arc::clone(&self.clients);
            tokio::spawn(async move {
                clients.write().await.remove(&target);
            });
            RPCError::Network(Self::network_error(e.to_string()))
        })?;

        let resp: AppendEntriesResponse<FlareNodeId> =
            serde_json::from_slice(&response.into_inner().data)
                .map_err(|e| RPCError::Network(Self::network_error(e.to_string())))?;

        Ok(resp)
    }

    async fn install_snapshot(
        &mut self,
        rpc: InstallSnapshotRequest<FlareTypeConfig>,
        _option: RPCOption,
    ) -> Result<InstallSnapshotResponse<FlareNodeId>, FlareInstallSnapshotError> {
        debug!(
            target = self.target,
            region = self.region_id,
            "sending install_snapshot"
        );

        let mut client = self
            .get_client()
            .await
            .map_err(|e| RPCError::Network(Self::network_error(e)))?;

        // Serialize the request
        let serialized = serde_json::to_vec(&rpc)
            .map_err(|e| RPCError::Network(Self::network_error(e.to_string())))?;

        // Send via gRPC
        let grpc_req = flaredb_proto::raft_server::OpenRaftSnapshotRequest {
            region_id: self.region_id,
            vote: vec![],
            meta: vec![],
            data: serialized,
        };

        let response = client.install_snapshot_v2(grpc_req).await.map_err(|e| {
            error!(error = %e, "install_snapshot RPC failed");
            RPCError::Network(Self::network_error(e.to_string()))
        })?;

        let resp: InstallSnapshotResponse<FlareNodeId> =
            serde_json::from_slice(&response.into_inner().data)
                .map_err(|e| RPCError::Network(Self::network_error(e.to_string())))?;

        Ok(resp)
    }

    async fn full_snapshot(
        &mut self,
        vote: Vote<FlareNodeId>,
        snapshot: Snapshot<FlareTypeConfig>,
        _cancel: impl Future<Output = ReplicationClosed> + OptionalSend + 'static,
        _option: RPCOption,
    ) -> Result<SnapshotResponse<FlareNodeId>, FlareSnapshotError> {
        debug!(
            target = self.target,
            region = self.region_id,
            snapshot_id = %snapshot.meta.snapshot_id,
            "sending full_snapshot"
        );

        let mut client = self
            .get_client()
            .await
            .map_err(|e| StreamingError::Network(Self::network_error(e)))?;

        // Read snapshot data
        let data = snapshot.snapshot.into_inner();

        // Serialize metadata
        let meta_bytes = serde_json::to_vec(&snapshot.meta)
            .map_err(|e| StreamingError::Network(Self::network_error(e.to_string())))?;

        let vote_bytes = serde_json::to_vec(&vote)
            .map_err(|e| StreamingError::Network(Self::network_error(e.to_string())))?;

        // Send via gRPC
        let grpc_req = flaredb_proto::raft_server::OpenRaftSnapshotRequest {
            region_id: self.region_id,
            vote: vote_bytes,
            meta: meta_bytes,
            data,
        };

        let response = client.install_snapshot_v2(grpc_req).await.map_err(|e| {
            error!(error = %e, "install_snapshot RPC failed");
            let target = self.target;
            let clients = Arc::clone(&self.clients);
            tokio::spawn(async move {
                clients.write().await.remove(&target);
            });
            StreamingError::Network(Self::network_error(e.to_string()))
        })?;

        let resp: SnapshotResponse<FlareNodeId> =
            serde_json::from_slice(&response.into_inner().data)
                .map_err(|e| StreamingError::Network(Self::network_error(e.to_string())))?;

        Ok(resp)
    }

    async fn vote(
        &mut self,
        rpc: VoteRequest<FlareNodeId>,
        _option: RPCOption,
    ) -> Result<VoteResponse<FlareNodeId>, FlareRPCError> {
        debug!(
            target = self.target,
            region = self.region_id,
            "sending vote request"
        );

        let mut client = self
            .get_client()
            .await
            .map_err(|e| RPCError::Network(Self::network_error(e)))?;

        // Serialize the request
        let serialized = serde_json::to_vec(&rpc)
            .map_err(|e| RPCError::Network(Self::network_error(e.to_string())))?;

        // Send via gRPC
        let grpc_req = flaredb_proto::raft_server::OpenRaftVoteRequest {
            region_id: self.region_id,
            data: serialized,
        };

        let response = client.vote_v2(grpc_req).await.map_err(|e| {
            error!(error = %e, "vote RPC failed");
            let target = self.target;
            let clients = Arc::clone(&self.clients);
            tokio::spawn(async move {
                clients.write().await.remove(&target);
            });
            RPCError::Network(Self::network_error(e.to_string()))
        })?;

        let resp: VoteResponse<FlareNodeId> = serde_json::from_slice(&response.into_inner().data)
            .map_err(|e| RPCError::Network(Self::network_error(e.to_string())))?;

        Ok(resp)
    }
}

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

    #[tokio::test]
    async fn test_network_factory_add_peer() {
        let factory = FlareNetworkFactory::new(1, 1);
        factory.add_peer(2, "127.0.0.1:50052".to_string()).await;

        let peers = factory.get_peers().await;
        assert_eq!(peers.get(&2), Some(&"127.0.0.1:50052".to_string()));
    }

    #[tokio::test]
    async fn test_network_factory_remove_peer() {
        let factory = FlareNetworkFactory::new(1, 1);
        factory.add_peer(2, "127.0.0.1:50052".to_string()).await;
        factory.remove_peer(2).await;

        let peers = factory.get_peers().await;
        assert!(peers.get(&2).is_none());
    }
}
