use std::sync::Arc;

use anyhow::Result;
use log::info;
use tici::worker::{
    WorkerNodeServiceCore,
    server::{DispatchShardWriter, WorkerServiceServer},
};
use tici_grpc::GrpcWorkerServiceClient;
use tici_shard::config::WorkerNodeConfig;
use tonic::transport::Channel;

pub struct WorkerServiceHandle {
    pub config: WorkerNodeConfig,
    pub server: WorkerServiceServer,
    pub client: GrpcWorkerServiceClient<Channel>,
    grpc_server_close_tx: Option<tokio::sync::oneshot::Sender<()>>,
    pub shard_write_manager: DispatchShardWriter,
    pub worker_node_service: Arc<WorkerNodeServiceCore>,
    pub data_dir_holder: Option<tempfile::TempDir>,
}

impl WorkerServiceHandle {
    fn prepare_data_dir(
        config: &mut WorkerNodeConfig,
        data_dir: Option<&str>,
    ) -> Result<Option<tempfile::TempDir>> {
        // If data_dir is not none, the caller is responsible for managing the data
        // dir's lifecycle.
        if let Some(data_dir) = data_dir {
            config.storage.data_dir = data_dir.to_string();
            Ok(None)
        } else {
            let tmp_dir = tempfile::tempdir()?;
            config.storage.data_dir = tmp_dir.path().to_str().unwrap().to_string();
            Ok(Some(tmp_dir))
        }
    }

    async fn start_server_and_build_handle(
        config: WorkerNodeConfig,
        data_dir_holder: Option<tempfile::TempDir>,
        mut server: WorkerServiceServer,
        grpc_server_close_tx: tokio::sync::oneshot::Sender<()>,
        grpc_server_close_rx: tokio::sync::oneshot::Receiver<()>,
    ) -> Result<Self> {
        let worker_node_service = server.worker_node_service.clone();
        let shard_write_manager = worker_node_service.shard_writer_manager();
        server.start(grpc_server_close_rx).await?;

        let client =
            WorkerNodeServiceCore::get_client(config.server.get_advertise_addr().to_string())
                .await?;
        info!("Worker server started at {}", config.server.addr);
        Ok(Self {
            config,
            server,
            client,
            grpc_server_close_tx: Some(grpc_server_close_tx),
            shard_write_manager,
            worker_node_service,
            data_dir_holder,
        })
    }

    pub async fn new_with_worker(
        mut config: WorkerNodeConfig,
        writer: DispatchShardWriter,
        data_dir: Option<&str>,
    ) -> Result<Self> {
        let data_dir_holder = Self::prepare_data_dir(&mut config, data_dir)?;
        let (tx, rx) = tokio::sync::oneshot::channel::<()>();
        let server = WorkerServiceServer::new_with_writer(config.clone(), writer).await?;
        info!(
            "new_with_writer: server created at {}, data_dir={}, data_dir_holder={:?}",
            config.server.addr, config.storage.data_dir, data_dir_holder
        );
        Self::start_server_and_build_handle(config, data_dir_holder, server, tx, rx).await
    }

    pub async fn new(mut config: WorkerNodeConfig, data_dir: Option<&str>) -> Result<Self> {
        info!("Creating worker service with config={:?}", config);
        let data_dir_holder = Self::prepare_data_dir(&mut config, data_dir)?;
        let (tx, rx) = tokio::sync::oneshot::channel::<()>();
        let server = WorkerServiceServer::new(config.clone()).await?;
        Self::start_server_and_build_handle(config, data_dir_holder, server, tx, rx).await
    }

    pub async fn stop(mut self) -> Result<()> {
        self.server
            .stop(self.grpc_server_close_tx.take().unwrap())
            .await?;
        info!("Worker server stopped");
        Ok(())
    }
}