// Copyright 2021 TiKV Project Authors. Licensed under Apache-2.0.
// Modifications Copyright 2022 PingCAP, Inc.

use std::sync::atomic::*;

use futures::{FutureExt, SinkExt, StreamExt, TryFutureExt, channel::mpsc, executor::block_on};
use grpcio::{self, *};
use kvproto::brpb::*;
use tikv_util::{error, info, worker::*};

use super::Task;

/// Service handles the RPC messages for the `Backup` service.

#[derive(Clone)]
pub struct Service {
    scheduler: Scheduler<Task>,
}

impl Service {
    /// Create a new backup service.
    pub fn new(scheduler: Scheduler<Task>) -> Self {
        Service { scheduler }
    }
}

impl Backup for Service {
    fn check_pending_admin_op(
        &mut self,
        ctx: RpcContext<'_>,
        _req: CheckAdminRequest,
        mut sink: ServerStreamingSink<CheckAdminResponse>,
    ) {
        let (mut tx, rx) = mpsc::unbounded();
        // self.router.broadcast_normal(|| {
        //     PeerMsg::SignificantMsg(SignificantMsg::CheckPendingAdmin(tx.clone()))
        // });
        block_on(tx.send(CheckAdminResponse::default())).unwrap();

        let send_task = async move {
            let mut s = rx.map(|resp| Ok((resp, WriteFlags::default())));
            sink.send_all(&mut s).await?;
            sink.close().await?;
            Ok(())
        }
        .map(|res: Result<()>| match res {
            Ok(_) => {
                info!("check admin closed");
            }
            Err(e) => {
                error!("check admin canceled"; "error" => ?e);
            }
        });
        ctx.spawn(send_task);
    }

    fn backup(
        &mut self,
        ctx: RpcContext<'_>,
        req: BackupRequest,
        mut sink: ServerStreamingSink<BackupResponse>,
    ) {
        let mut cancel = None;
        // TODO: make it a bounded channel.
        let (tx, rx) = mpsc::unbounded();
        if let Err(status) = match Task::new(req, tx) {
            Ok((task, c)) => {
                cancel = Some(c);
                self.scheduler.schedule(task).map_err(|e| {
                    RpcStatus::with_message(RpcStatusCode::INVALID_ARGUMENT, format!("{:?}", e))
                })
            }
            Err(e) => Err(RpcStatus::with_message(
                RpcStatusCode::UNKNOWN,
                format!("{:?}", e),
            )),
        } {
            error!("backup task initiate failed"; "error" => ?status);
            ctx.spawn(
                sink.fail(status)
                    .unwrap_or_else(|e| error!("backup failed to send error"; "error" => ?e)),
            );
            return;
        };

        let send_task = async move {
            let mut s = rx.map(|resp| Ok((resp, WriteFlags::default())));
            sink.send_all(&mut s).await?;
            sink.close().await?;
            Ok(())
        }
        .map(|res: Result<()>| {
            match res {
                Ok(_) => {
                    info!("backup closed");
                }
                Err(e) => {
                    if let Some(c) = cancel {
                        // Cancel the running task.
                        c.store(true, Ordering::SeqCst);
                    }
                    error!("backup canceled"; "error" => ?e);
                }
            }
        });

        ctx.spawn(send_task);
    }
}