//! Node client for communicating with storage nodes

use super::{NodeError, NodeResult};
use async_trait::async_trait;
use bytes::Bytes;
use futures::stream;
use lightningstor_node::proto::{
    BatchGetChunksRequest, ChunkExistsRequest, ChunkSizeRequest, DeleteChunkRequest,
    GetChunkRequest, PingRequest, PutChunkRequest,
};
use lightningstor_node::NodeServiceClient;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::RwLock;
use tokio::time::timeout;
use tonic::transport::{Channel, Endpoint};

const MAX_NODE_GRPC_MESSAGE_SIZE: usize = 1024 * 1024 * 1024;
const DEFAULT_NODE_CONNECT_TIMEOUT: Duration = Duration::from_secs(5);
const DEFAULT_NODE_RPC_TIMEOUT: Duration = Duration::from_secs(300);
const NODE_HEALTH_PROBE_TIMEOUT: Duration = Duration::from_secs(5);
const NODE_GRPC_INITIAL_STREAM_WINDOW: u32 = 64 * 1024 * 1024;
const NODE_GRPC_INITIAL_CONNECTION_WINDOW: u32 = 512 * 1024 * 1024;
const NODE_GRPC_KEEPALIVE_INTERVAL: Duration = Duration::from_secs(30);
const NODE_GRPC_KEEPALIVE_TIMEOUT: Duration = Duration::from_secs(10);
const TRANSIENT_RPC_RETRY_ATTEMPTS: usize = 2;

/// Trait for storage node client operations
#[async_trait]
pub trait NodeClientTrait: Send + Sync {
    /// Get the node ID
    fn node_id(&self) -> &str;

    /// Get the node endpoint
    fn endpoint(&self) -> &str;

    /// Check if the node is currently considered healthy
    async fn is_healthy(&self) -> bool;

    /// Store a chunk on this node
    async fn put_chunk(
        &self,
        chunk_id: &str,
        shard_index: u32,
        is_parity: bool,
        data: Bytes,
    ) -> NodeResult<()>;

    /// Store multiple chunks on this node using a more efficient batch path when available.
    async fn batch_put_chunks(
        &self,
        chunks: Vec<(String, u32, bool, Bytes)>,
    ) -> NodeResult<()> {
        for (chunk_id, shard_index, is_parity, data) in chunks {
            self.put_chunk(&chunk_id, shard_index, is_parity, data).await?;
        }
        Ok(())
    }

    /// Retrieve a chunk from this node
    async fn get_chunk(
        &self,
        chunk_id: &str,
        shard_index: u32,
        is_parity: bool,
    ) -> NodeResult<Bytes>;

    /// Retrieve multiple chunks from this node using a more efficient batch path when available.
    async fn batch_get_chunks(&self, requests: Vec<(String, u32, bool)>) -> NodeResult<Vec<Bytes>> {
        let mut chunks = Vec::with_capacity(requests.len());
        for (chunk_id, shard_index, is_parity) in requests {
            chunks.push(self.get_chunk(&chunk_id, shard_index, is_parity).await?);
        }
        Ok(chunks)
    }

    /// Delete a chunk from this node
    async fn delete_chunk(&self, chunk_id: &str) -> NodeResult<()>;

    /// Check if a chunk exists on this node
    async fn chunk_exists(&self, chunk_id: &str) -> NodeResult<bool>;

    /// Get the size of a chunk on this node
    async fn chunk_size(&self, chunk_id: &str) -> NodeResult<Option<u64>>;

    /// Ping the node to check connectivity
    async fn ping(&self) -> NodeResult<Duration>;
}

/// Real gRPC client for storage nodes
///
/// This client communicates with storage nodes over gRPC.
/// For now, this is a placeholder that will be implemented
/// when the storage node service is created.
pub struct NodeClient {
    node_id: String,
    endpoint: String,
    healthy: AtomicBool,
    connection_timeout: Duration,
    request_timeout: Duration,
    client: RwLock<NodeServiceClient<Channel>>,
}

impl NodeClient {
    fn endpoint_builder(
        endpoint_url: String,
        connection_timeout: Duration,
    ) -> NodeResult<Endpoint> {
        Endpoint::from_shared(endpoint_url.clone())
            .map_err(|e| NodeError::ConnectionFailed {
                node_id: "unknown".to_string(),
                reason: e.to_string(),
            })
            .map(|endpoint| {
                endpoint
                    .connect_timeout(connection_timeout)
                    .tcp_nodelay(true)
                    .http2_keep_alive_interval(NODE_GRPC_KEEPALIVE_INTERVAL)
                    .keep_alive_timeout(NODE_GRPC_KEEPALIVE_TIMEOUT)
                    .initial_stream_window_size(NODE_GRPC_INITIAL_STREAM_WINDOW)
                    .initial_connection_window_size(NODE_GRPC_INITIAL_CONNECTION_WINDOW)
            })
    }

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

    fn client_from_channel(channel: Channel) -> NodeServiceClient<Channel> {
        NodeServiceClient::new(channel)
            .max_decoding_message_size(MAX_NODE_GRPC_MESSAGE_SIZE)
            .max_encoding_message_size(MAX_NODE_GRPC_MESSAGE_SIZE)
    }

    async fn build_client(
        endpoint: &str,
        connection_timeout: Duration,
        lazy: bool,
    ) -> NodeResult<NodeServiceClient<Channel>> {
        let endpoint_url = Self::normalize_endpoint(endpoint);
        let endpoint = Self::endpoint_builder(endpoint_url, connection_timeout)?;
        let channel = if lazy {
            endpoint.connect_lazy()
        } else {
            endpoint
                .connect()
                .await
                .map_err(|e| NodeError::ConnectionFailed {
                    node_id: "unknown".to_string(),
                    reason: e.to_string(),
                })?
        };
        Ok(Self::client_from_channel(channel))
    }

    /// Connect to a storage node at the given endpoint
    pub async fn connect(endpoint: &str) -> NodeResult<Self> {
        Self::connect_with_timeouts(
            endpoint,
            DEFAULT_NODE_CONNECT_TIMEOUT,
            DEFAULT_NODE_RPC_TIMEOUT,
        )
        .await
    }

    /// Connect to a storage node with explicit timeouts.
    pub async fn connect_with_timeouts(
        endpoint: &str,
        connection_timeout: Duration,
        request_timeout: Duration,
    ) -> NodeResult<Self> {
        let client = Self::build_client(endpoint, connection_timeout, false).await?;

        // Try to get node status to get the real node ID
        // If that fails, generate a temporary one based on endpoint, but connection is established
        let node_id = match client
            .clone()
            .get_status(lightningstor_node::proto::GetStatusRequest {})
            .await
        {
            Ok(response) => response.into_inner().node_id,
            Err(_) => format!("node-{}", endpoint.replace([':', '.', '/'], "-")),
        };

        Ok(Self {
            node_id,
            endpoint: endpoint.to_string(),
            healthy: AtomicBool::new(true),
            connection_timeout,
            request_timeout,
            client: RwLock::new(client),
        })
    }

    /// Create a client with a specific node ID
    pub async fn connect_with_id(node_id: &str, endpoint: &str) -> NodeResult<Self> {
        Self::connect_with_id_and_timeouts(
            node_id,
            endpoint,
            DEFAULT_NODE_CONNECT_TIMEOUT,
            DEFAULT_NODE_RPC_TIMEOUT,
        )
        .await
    }

    /// Create a client with a specific node ID and explicit timeouts.
    pub async fn connect_with_id_and_timeouts(
        node_id: &str,
        endpoint: &str,
        connection_timeout: Duration,
        request_timeout: Duration,
    ) -> NodeResult<Self> {
        let client = Self::build_client(endpoint, connection_timeout, true)
            .await
            .map_err(|e| match e {
                NodeError::ConnectionFailed { reason, .. } => NodeError::ConnectionFailed {
                    node_id: node_id.to_string(),
                    reason,
                },
                other => other,
            })?;

        Ok(Self {
            node_id: node_id.to_string(),
            endpoint: endpoint.to_string(),
            healthy: AtomicBool::new(true),
            connection_timeout,
            request_timeout,
            client: RwLock::new(client),
        })
    }

    /// Mark the node as unhealthy
    pub fn mark_unhealthy(&self) {
        self.healthy.store(false, Ordering::SeqCst);
    }

    /// Mark the node as healthy
    pub fn mark_healthy(&self) {
        self.healthy.store(true, Ordering::SeqCst);
    }

    async fn clone_client(&self) -> NodeServiceClient<Channel> {
        self.client.read().await.clone()
    }

    async fn reconnect(&self) -> NodeResult<()> {
        let client = Self::build_client(&self.endpoint, self.connection_timeout, true)
            .await
            .map_err(|e| match e {
                NodeError::ConnectionFailed { reason, .. } => NodeError::ConnectionFailed {
                    node_id: self.node_id.clone(),
                    reason,
                },
                other => other,
            })?;
        *self.client.write().await = client;
        Ok(())
    }

    fn is_retryable_status(status: &tonic::Status) -> bool {
        matches!(
            status.code(),
            tonic::Code::Unavailable
                | tonic::Code::Cancelled
                | tonic::Code::Unknown
                | tonic::Code::DeadlineExceeded
                | tonic::Code::ResourceExhausted
        ) || status.message().contains("transport error")
            || status.message().contains("http2 error")
    }

    fn health_probe_timeout(&self) -> Duration {
        self.request_timeout.min(NODE_HEALTH_PROBE_TIMEOUT)
    }

    async fn probe_health(&self) -> bool {
        for attempt in 0..=1 {
            let mut client = self.clone_client().await;
            match timeout(self.health_probe_timeout(), client.ping(PingRequest {})).await {
                Ok(Ok(_)) => {
                    self.mark_healthy();
                    return true;
                }
                Ok(Err(status)) if attempt == 0 && Self::is_retryable_status(&status) => {
                    if self.reconnect().await.is_ok() {
                        continue;
                    }
                }
                Ok(Err(_)) | Err(_) => {}
            }
        }
        self.mark_unhealthy();
        false
    }
}

#[async_trait]
impl NodeClientTrait for NodeClient {
    fn node_id(&self) -> &str {
        &self.node_id
    }

    fn endpoint(&self) -> &str {
        &self.endpoint
    }

    async fn is_healthy(&self) -> bool {
        if self.healthy.load(Ordering::SeqCst) {
            return true;
        }

        self.probe_health().await
    }

    async fn put_chunk(
        &self,
        chunk_id: &str,
        shard_index: u32,
        is_parity: bool,
        data: Bytes,
    ) -> NodeResult<()> {
        if !self.is_healthy().await {
            return Err(NodeError::Unhealthy(self.node_id.clone()));
        }

        for attempt in 0..TRANSIENT_RPC_RETRY_ATTEMPTS {
            let request = PutChunkRequest {
                chunk_id: chunk_id.to_string(),
                shard_index,
                is_parity,
                data: data.clone(),
            };
            let mut client = self.clone_client().await;
            let result = timeout(self.request_timeout, client.put_chunk(request)).await;
            match result {
                Ok(Ok(_)) => {
                    self.mark_healthy();
                    return Ok(());
                }
                Ok(Err(status))
                    if attempt + 1 < TRANSIENT_RPC_RETRY_ATTEMPTS
                        && Self::is_retryable_status(&status) =>
                {
                    self.mark_unhealthy();
                    let _ = self.reconnect().await;
                }
                Ok(Err(status)) => {
                    self.mark_unhealthy();
                    return Err(NodeError::RpcFailed(status.to_string()));
                }
                Err(_) if attempt + 1 < TRANSIENT_RPC_RETRY_ATTEMPTS => {
                    self.mark_unhealthy();
                    let _ = self.reconnect().await;
                }
                Err(_) => {
                    self.mark_unhealthy();
                    return Err(NodeError::Timeout);
                }
            }
        }
        self.mark_unhealthy();
        Err(NodeError::Timeout)
    }

    async fn batch_put_chunks(
        &self,
        chunks: Vec<(String, u32, bool, Bytes)>,
    ) -> NodeResult<()> {
        if !self.is_healthy().await {
            return Err(NodeError::Unhealthy(self.node_id.clone()));
        }

        let request_count = chunks.len();

        for attempt in 0..TRANSIENT_RPC_RETRY_ATTEMPTS {
            let request_stream = stream::iter(chunks.clone().into_iter().map(
                |(chunk_id, shard_index, is_parity, data)| PutChunkRequest {
                    chunk_id,
                    shard_index,
                    is_parity,
                    data,
                },
            ));

            let mut client = self.clone_client().await;
            match timeout(self.request_timeout, client.batch_put_chunks(request_stream)).await {
                Ok(Ok(response)) => {
                    let response = response.into_inner();
                    if response.failure_count == 0 && response.success_count as usize == request_count
                    {
                        self.mark_healthy();
                        return Ok(());
                    }
                    self.mark_unhealthy();
                    return Err(NodeError::RpcFailed(format!(
                        "batch put stored {} of {} chunks: {}",
                        response.success_count,
                        request_count,
                        response.errors.join("; ")
                    )));
                }
                Ok(Err(status))
                    if attempt + 1 < TRANSIENT_RPC_RETRY_ATTEMPTS
                        && Self::is_retryable_status(&status) =>
                {
                    self.mark_unhealthy();
                    let _ = self.reconnect().await;
                }
                Ok(Err(status)) => {
                    self.mark_unhealthy();
                    return Err(NodeError::RpcFailed(status.to_string()));
                }
                Err(_) if attempt + 1 < TRANSIENT_RPC_RETRY_ATTEMPTS => {
                    self.mark_unhealthy();
                    let _ = self.reconnect().await;
                }
                Err(_) => {
                    self.mark_unhealthy();
                    return Err(NodeError::Timeout);
                }
            }
        }
        self.mark_unhealthy();
        Err(NodeError::Timeout)
    }

    async fn get_chunk(
        &self,
        chunk_id: &str,
        shard_index: u32,
        is_parity: bool,
    ) -> NodeResult<Bytes> {
        if !self.is_healthy().await {
            return Err(NodeError::Unhealthy(self.node_id.clone()));
        }

        let request = GetChunkRequest {
            chunk_id: chunk_id.to_string(),
            shard_index,
            is_parity,
        };

        let mut client = self.clone_client().await;
        let response = match timeout(self.request_timeout, client.get_chunk(request)).await {
            Ok(Ok(response)) => {
                self.mark_healthy();
                response
            }
            Ok(Err(e)) => match e.code() {
                tonic::Code::NotFound => {
                    self.mark_healthy();
                    return Err(NodeError::NotFound(chunk_id.to_string()));
                }
                _ => {
                    self.mark_unhealthy();
                    return Err(NodeError::RpcFailed(e.to_string()));
                }
            },
            Err(_) => {
                self.mark_unhealthy();
                return Err(NodeError::Timeout);
            }
        };

        Ok(response.into_inner().data)
    }

    async fn batch_get_chunks(&self, requests: Vec<(String, u32, bool)>) -> NodeResult<Vec<Bytes>> {
        if !self.is_healthy().await {
            return Err(NodeError::Unhealthy(self.node_id.clone()));
        }

        let request_count = requests.len();
        let request = BatchGetChunksRequest {
            chunks: requests
                .into_iter()
                .map(|(chunk_id, shard_index, is_parity)| GetChunkRequest {
                    chunk_id,
                    shard_index,
                    is_parity,
                })
                .collect(),
        };

        let mut client = self.clone_client().await;
        let response = match timeout(self.request_timeout, client.batch_get_chunks(request)).await {
            Ok(Ok(response)) => {
                self.mark_healthy();
                response
            }
            Ok(Err(e)) => {
                self.mark_unhealthy();
                return Err(NodeError::RpcFailed(e.to_string()));
            }
            Err(_) => {
                self.mark_unhealthy();
                return Err(NodeError::Timeout);
            }
        };

        let mut stream = response.into_inner();
        let mut chunks = Vec::with_capacity(request_count);
        loop {
            let next = match timeout(self.request_timeout, stream.message()).await {
                Ok(Ok(next)) => next,
                Ok(Err(e)) => {
                    self.mark_unhealthy();
                    return match e.code() {
                        tonic::Code::NotFound => Err(NodeError::NotFound("batch-get".into())),
                        _ => Err(NodeError::RpcFailed(e.to_string())),
                    };
                }
                Err(_) => {
                    self.mark_unhealthy();
                    return Err(NodeError::Timeout);
                }
            };

            let Some(chunk) = next else {
                break;
            };
            chunks.push(chunk.data);
        }

        if chunks.len() != request_count {
            self.mark_unhealthy();
            return Err(NodeError::RpcFailed(format!(
                "batch get returned {} of {} chunks",
                chunks.len(),
                request_count
            )));
        }

        self.mark_healthy();
        Ok(chunks)
    }

    async fn delete_chunk(&self, chunk_id: &str) -> NodeResult<()> {
        if !self.is_healthy().await {
            return Err(NodeError::Unhealthy(self.node_id.clone()));
        }

        let request = DeleteChunkRequest {
            chunk_id: chunk_id.to_string(),
        };

        let mut client = self.clone_client().await;
        let result = timeout(self.request_timeout, client.delete_chunk(request)).await;
        match result {
            Ok(Ok(_)) => {
                self.mark_healthy();
                Ok(())
            }
            Ok(Err(e)) => {
                self.mark_unhealthy();
                Err(NodeError::RpcFailed(e.to_string()))
            }
            Err(_) => {
                self.mark_unhealthy();
                Err(NodeError::Timeout)
            }
        }
    }

    async fn chunk_exists(&self, chunk_id: &str) -> NodeResult<bool> {
        if !self.is_healthy().await {
            return Err(NodeError::Unhealthy(self.node_id.clone()));
        }

        let request = ChunkExistsRequest {
            chunk_id: chunk_id.to_string(),
        };

        let mut client = self.clone_client().await;
        let response = match timeout(self.request_timeout, client.chunk_exists(request)).await {
            Ok(Ok(response)) => {
                self.mark_healthy();
                response
            }
            Ok(Err(e)) => {
                self.mark_unhealthy();
                return Err(NodeError::RpcFailed(e.to_string()));
            }
            Err(_) => {
                self.mark_unhealthy();
                return Err(NodeError::Timeout);
            }
        };

        Ok(response.into_inner().exists)
    }

    async fn chunk_size(&self, chunk_id: &str) -> NodeResult<Option<u64>> {
        if !self.is_healthy().await {
            return Err(NodeError::Unhealthy(self.node_id.clone()));
        }

        let request = ChunkSizeRequest {
            chunk_id: chunk_id.to_string(),
        };

        let mut client = self.clone_client().await;
        let response = match timeout(self.request_timeout, client.chunk_size(request)).await {
            Ok(Ok(response)) => {
                self.mark_healthy();
                response
            }
            Ok(Err(e)) => {
                self.mark_unhealthy();
                return Err(NodeError::RpcFailed(e.to_string()));
            }
            Err(_) => {
                self.mark_unhealthy();
                return Err(NodeError::Timeout);
            }
        };

        let inner = response.into_inner();
        if inner.exists {
            Ok(Some(inner.size))
        } else {
            Ok(None)
        }
    }

    async fn ping(&self) -> NodeResult<Duration> {
        if !self.is_healthy().await {
            return Err(NodeError::Unhealthy(self.node_id.clone()));
        }

        let start = std::time::Instant::now();
        let request = PingRequest {};

        let mut client = self.clone_client().await;
        match timeout(self.request_timeout, client.ping(request)).await {
            Ok(Ok(_)) => {
                self.mark_healthy();
            }
            Ok(Err(e)) => {
                self.mark_unhealthy();
                return Err(NodeError::RpcFailed(e.to_string()));
            }
            Err(_) => {
                self.mark_unhealthy();
                return Err(NodeError::Timeout);
            }
        }

        Ok(start.elapsed())
    }
}

/// A pool of node clients for connection reuse
pub struct NodeClientPool {
    clients: RwLock<Vec<Arc<dyn NodeClientTrait>>>,
}

impl NodeClientPool {
    /// Create a new empty client pool
    pub fn new() -> Self {
        Self {
            clients: RwLock::new(Vec::new()),
        }
    }

    /// Add a client to the pool
    pub async fn add(&self, client: Arc<dyn NodeClientTrait>) {
        self.clients.write().await.push(client);
    }

    /// Get all clients in the pool
    pub async fn all(&self) -> Vec<Arc<dyn NodeClientTrait>> {
        self.clients.read().await.clone()
    }

    /// Get all healthy clients
    pub async fn healthy(&self) -> Vec<Arc<dyn NodeClientTrait>> {
        let clients = self.clients.read().await;
        let mut healthy = Vec::new();
        for client in clients.iter() {
            if client.is_healthy().await {
                healthy.push(client.clone());
            }
        }
        healthy
    }

    /// Get a client by node ID
    pub async fn get(&self, node_id: &str) -> Option<Arc<dyn NodeClientTrait>> {
        self.clients
            .read()
            .await
            .iter()
            .find(|c| c.node_id() == node_id)
            .cloned()
    }

    /// Remove a client from the pool
    pub async fn remove(&self, node_id: &str) {
        self.clients
            .write()
            .await
            .retain(|c| c.node_id() != node_id);
    }

    /// Get the number of clients in the pool
    pub async fn len(&self) -> usize {
        self.clients.read().await.len()
    }

    /// Check if the pool is empty
    pub async fn is_empty(&self) -> bool {
        self.clients.read().await.is_empty()
    }
}

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

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

    #[tokio::test]
    async fn test_node_client_creation() {
        let client = NodeClient::connect_with_id("test-node", "http://localhost:9002")
            .await
            .unwrap();
        assert!(client.is_healthy().await);
        assert_eq!(client.node_id(), "test-node");
    }

    #[tokio::test]
    async fn test_node_client_health_toggle() {
        let client = NodeClient::connect_with_id("test-node", "http://localhost:9002")
            .await
            .unwrap();

        assert!(client.is_healthy().await);
        client.mark_unhealthy();
        assert!(!client.is_healthy().await);
        client.mark_healthy();
        assert!(client.is_healthy().await);
    }

    #[tokio::test]
    async fn test_node_client_pool() {
        let pool = NodeClientPool::new();
        assert!(pool.is_empty().await);

        let client1 = Arc::new(
            NodeClient::connect_with_id("node1", "http://node1:9002")
                .await
                .unwrap(),
        );
        let client2 = Arc::new(
            NodeClient::connect_with_id("node2", "http://node2:9002")
                .await
                .unwrap(),
        );

        pool.add(client1.clone()).await;
        pool.add(client2.clone()).await;

        assert_eq!(pool.len().await, 2);
        assert!(pool.get(client1.node_id()).await.is_some());

        pool.remove(client1.node_id()).await;
        assert_eq!(pool.len().await, 1);
        assert!(pool.get(client1.node_id()).await.is_none());
    }
}
