//! Lease service implementation

use crate::conversions::make_header;
use crate::proto::{
    lease_server::Lease, LeaseGrantRequest, LeaseGrantResponse, LeaseKeepAliveRequest,
    LeaseKeepAliveResponse, LeaseLeasesRequest, LeaseLeasesResponse, LeaseRevokeRequest,
    LeaseRevokeResponse, LeaseStatus, LeaseTimeToLiveRequest, LeaseTimeToLiveResponse,
};
use chainfire_raft::core::RaftCore;
use chainfire_types::command::RaftCommand;
use std::pin::Pin;
use std::sync::Arc;
use tokio::sync::mpsc;
use tokio_stream::{wrappers::ReceiverStream, Stream, StreamExt};
use tonic::{Request, Response, Status, Streaming};
use tracing::{debug, warn};

/// Lease service implementation
pub struct LeaseServiceImpl {
    /// Raft core for consensus
    raft: Arc<RaftCore>,
    /// Cluster ID
    cluster_id: u64,
}

impl LeaseServiceImpl {
    /// Create a new Lease service
    pub fn new(raft: Arc<RaftCore>, cluster_id: u64) -> Self {
        Self { raft, cluster_id }
    }

    /// Create a response header
    fn make_header(&self, revision: u64) -> crate::proto::ResponseHeader {
        make_header(self.cluster_id, self.raft.id(), revision, 0)
    }
}

#[tonic::async_trait]
impl Lease for LeaseServiceImpl {
    async fn lease_grant(
        &self,
        request: Request<LeaseGrantRequest>,
    ) -> Result<Response<LeaseGrantResponse>, Status> {
        let req = request.into_inner();
        debug!(id = req.id, ttl = req.ttl, "LeaseGrant request");

        let command = RaftCommand::LeaseGrant {
            id: req.id,
            ttl: req.ttl,
        };

        let response = self
            .raft
            .write(command)
            .await
            .map_err(|e| Status::internal(e.to_string()))?;

        Ok(Response::new(LeaseGrantResponse {
            header: Some(self.make_header(response.revision)),
            id: response.lease_id.unwrap_or(0),
            ttl: response.lease_ttl.unwrap_or(0),
            error: String::new(),
        }))
    }

    async fn lease_revoke(
        &self,
        request: Request<LeaseRevokeRequest>,
    ) -> Result<Response<LeaseRevokeResponse>, Status> {
        let req = request.into_inner();
        debug!(id = req.id, "LeaseRevoke request");

        let command = RaftCommand::LeaseRevoke { id: req.id };

        let response = self
            .raft
            .write(command)
            .await
            .map_err(|e| Status::internal(e.to_string()))?;

        Ok(Response::new(LeaseRevokeResponse {
            header: Some(self.make_header(response.revision)),
        }))
    }

    type LeaseKeepAliveStream =
        Pin<Box<dyn Stream<Item = Result<LeaseKeepAliveResponse, Status>> + Send>>;

    async fn lease_keep_alive(
        &self,
        request: Request<Streaming<LeaseKeepAliveRequest>>,
    ) -> Result<Response<Self::LeaseKeepAliveStream>, Status> {
        let mut stream = request.into_inner();
        let raft = Arc::clone(&self.raft);
        let cluster_id = self.cluster_id;

        let (tx, rx) = mpsc::channel(16);

        tokio::spawn(async move {
            while let Some(result) = stream.next().await {
                match result {
                    Ok(req) => {
                        debug!(id = req.id, "LeaseKeepAlive request");

                        let command = RaftCommand::LeaseRefresh { id: req.id };

                        match raft.write(command).await {
                            Ok(response) => {
                                let resp = LeaseKeepAliveResponse {
                                    header: Some(make_header(
                                        cluster_id,
                                        raft.id(),
                                        response.revision,
                                        0,
                                    )),
                                    id: response.lease_id.unwrap_or(req.id),
                                    ttl: response.lease_ttl.unwrap_or(0),
                                };
                                if tx.send(Ok(resp)).await.is_err() {
                                    break;
                                }
                            }
                            Err(e) => {
                                warn!("LeaseKeepAlive failed: {}", e);
                                if tx.send(Err(Status::internal(e.to_string()))).await.is_err() {
                                    break;
                                }
                            }
                        }
                    }
                    Err(e) => {
                        warn!("LeaseKeepAlive stream error: {}", e);
                        break;
                    }
                }
            }
        });

        Ok(Response::new(Box::pin(ReceiverStream::new(rx))))
    }

    async fn lease_time_to_live(
        &self,
        request: Request<LeaseTimeToLiveRequest>,
    ) -> Result<Response<LeaseTimeToLiveResponse>, Status> {
        let req = request.into_inner();
        debug!(id = req.id, "LeaseTimeToLive request");

        // Read directly from state machine
        let sm = self.raft.state_machine();
        let revision = sm.current_revision();

        let leases = sm.leases();
        match leases.time_to_live(req.id) {
            Some((ttl, granted_ttl, keys)) => Ok(Response::new(LeaseTimeToLiveResponse {
                header: Some(self.make_header(revision)),
                id: req.id,
                ttl,
                granted_ttl,
                keys: if req.keys { keys } else { vec![] },
            })),
            None => Ok(Response::new(LeaseTimeToLiveResponse {
                header: Some(self.make_header(revision)),
                id: req.id,
                ttl: -1,
                granted_ttl: 0,
                keys: vec![],
            })),
        }
    }

    async fn lease_leases(
        &self,
        _request: Request<LeaseLeasesRequest>,
    ) -> Result<Response<LeaseLeasesResponse>, Status> {
        debug!("LeaseLeases request");

        // Read directly from state machine
        let sm = self.raft.state_machine();
        let revision = sm.current_revision();

        let leases = sm.leases();
        let lease_ids = leases.list();

        let statuses: Vec<LeaseStatus> = lease_ids.into_iter().map(|id| LeaseStatus { id }).collect();

        Ok(Response::new(LeaseLeasesResponse {
            header: Some(self.make_header(revision)),
            leases: statuses,
        }))
    }
}
