use std::{net::SocketAddr, sync::Arc};

use anyhow::{Context, Result};
use async_trait::async_trait;
use log::{error, info, warn};
use storage_backend::storage::{STORAGE_SERVICE_WORKER, create_backend_storage};
use tici_common::config::ServerConfig;
use tici_grpc::{
    AddShardRequest, AddShardResponse, CompactFragRequest, CompactFragResponse, GrpcWorkerService,
    GrpcWorkerServiceServer, ImportWriteRequest, ImportWriteResponse, WorkerNodeShardStatus,
    metrics::begin_grpc_handle_metrics,
};
use tici_shard::{
    config::WorkerNodeConfig,
    metadata::metadata_store::MetadataStore,
    utils::{
        frag_path_generator::{FragPathGenerator, TSOGenerator},
        meta_client_pool::{GrpcMetaClientPool, MetaClientPool},
    },
    writer::{frag_batch_uploader::FragBatchUploader, shard_writer_manager::ShardWriterManager},
};
use tonic::{Request, Response, Status, transport::Server};

use super::WorkerNodeServiceCore;
use crate::{
    meta::writer::CompactionStateKey,
    status_server::{StatusServer, operator::Operators},
    worker::{mocks::MockShardWriterManager, service::AddShardSource},
};

#[derive(Clone)]
pub enum DispatchShardWriter {
    Mock(Arc<MockShardWriterManager>),
    Real(Arc<ShardWriterManager>),
}

impl DispatchShardWriter {
    pub fn is_mocked(&self) -> bool {
        match self {
            Self::Mock(_) => true,
            Self::Real(_) => false,
        }
    }

    pub fn shard_statuses(&self) -> Vec<WorkerNodeShardStatus> {
        match self {
            Self::Mock(mock) => mock.shard_statuses(),
            Self::Real(manager) => manager.shard_statuses(),
        }
    }

    #[cfg(any(test, feature = "testexport"))]
    pub fn as_mock(&self) -> Arc<MockShardWriterManager> {
        match self {
            Self::Mock(mock) => mock.clone(),
            Self::Real(_) => panic!("must be mock"),
        }
    }

    pub async fn frag_batch_uploader(&self) -> Option<Arc<FragBatchUploader>> {
        match self {
            Self::Real(manager) => Some(manager.frag_batch_uploader()),
            Self::Mock(_) => None,
        }
    }
}

#[derive(Clone)]
pub struct WorkerGrpcServiceImpl {
    worker_node_service: Arc<WorkerNodeServiceCore>,
}

impl WorkerGrpcServiceImpl {
    pub fn new(worker_node_service: Arc<WorkerNodeServiceCore>) -> Self {
        Self {
            worker_node_service,
        }
    }
}

#[async_trait]
impl GrpcWorkerService for WorkerGrpcServiceImpl {
    async fn add_shard(
        &self,
        request: Request<AddShardRequest>,
    ) -> Result<Response<AddShardResponse>, Status> {
        let rpc_metrics = begin_grpc_handle_metrics("worker", "add_shard");
        fail::fail_point!("worker_service_mock_drop_add_shard", |_| {
            Ok(Response::new(AddShardResponse { status: 0 }))
        });
        let req = request.into_inner();
        let (shard_meta, index_meta) =
            WorkerNodeServiceCore::parse_add_shard_req(&req).map_err(|e| {
                rpc_metrics.mark_error();
                warn!("add_shard: failed to parse req, err={:?}", e);
                Status::internal(format!("Failed to parse req, err={}", e))
            })?;
        match self
            .worker_node_service
            .add_shard(shard_meta, index_meta, AddShardSource::Rpc)
            .await
        {
            Ok(_) => Ok(Response::new(AddShardResponse { status: 0 })),
            Err(e) => {
                rpc_metrics.mark_error();
                warn!("add_shard={:?}", e);
                Ok(Response::new(AddShardResponse { status: 1 }))
            }
        }
    }

    async fn compact_fragments(
        &self,
        request: Request<CompactFragRequest>,
    ) -> Result<Response<CompactFragResponse>, Status> {
        let rpc_metrics = begin_grpc_handle_metrics("worker", "compact_fragments");
        fail::fail_point!("worker_service_mock_drop_compaction", |_| {
            Ok(Response::new(CompactFragResponse { status: 0 }))
        });
        let req = request.into_inner();
        let compaction_key = CompactionStateKey::try_from(&req).map_err(|e| {
            rpc_metrics.mark_error();
            warn!(
                "compact_fragments: failed to create compaction key, err={:?}",
                e
            );
            Status::internal(format!("Failed to create compaction key, err={}", e))
        })?;

        let output_frag_path = self
            .worker_node_service
            .frag_path_gen
            .new_path(&req.get_index_key())
            .await
            .map_err(|e| {
                rpc_metrics.mark_error();
                error!(
                    "compact_fragments: failed to generate frag path, key={:?}, err={:?}",
                    compaction_key, e
                );
                Status::internal(format!("Failed to generate frag path, err={}", e))
            })?;

        match self
            .worker_node_service
            .compact_fragments(compaction_key.clone(), req.clone(), output_frag_path)
            .await
        {
            Ok(()) => {
                info!(
                    "compact_fragments: scheduled successfully, key={:?}, frags={:?}",
                    compaction_key,
                    req.input_frags
                        .iter()
                        .map(|f| &f.frag_path)
                        .collect::<Vec<_>>()
                );
                Ok(Response::new(CompactFragResponse { status: 0 }))
            }
            Err(e) => {
                rpc_metrics.mark_error();
                warn!(
                    "compact_fragments: schedule failed, key={:?}, err={:?}",
                    compaction_key, e
                );
                Err(Status::internal(format!(
                    "Failed to schedule compaction, err={}",
                    e
                )))
            }
        }
    }

    async fn import_write(
        &self,
        request: Request<ImportWriteRequest>,
    ) -> Result<Response<ImportWriteResponse>, Status> {
        let rpc_metrics = begin_grpc_handle_metrics("worker", "import_write");
        let req = request.into_inner();
        if req.file_location.is_empty() {
            rpc_metrics.mark_error();
            return Err(Status::invalid_argument(
                "Missing import file location in ImportWriteRequest",
            ));
        }
        match self.worker_node_service.import_write(req.clone()).await {
            Ok(()) => Ok(Response::new(ImportWriteResponse { status: 0 })),
            Err(e) => {
                rpc_metrics.mark_error();
                warn!("import_write call return, req={:?}, err={}", req, e);
                Ok(Response::new(ImportWriteResponse { status: 1 }))
            }
        }
    }
}

/// WorkerServiceServer controls the lifetime of the whole service
pub struct WorkerServiceServer {
    status_server: Option<StatusServer>,
    pub worker_node_service: Arc<WorkerNodeServiceCore>,
    pub config: WorkerNodeConfig,
    grpc_handle: Option<tokio::task::JoinHandle<Result<(), tonic::transport::Error>>>,
}

impl WorkerServiceServer {
    pub async fn new(config: WorkerNodeConfig) -> Result<Self> {
        config.storage.init_worker_data_dir()?;
        let frag_path_gen = Self::create_frag_path_gen(&config).await?;
        let pd_addrs = config.server.get_pd_addresses();
        let meta_client_pool: Arc<dyn MetaClientPool> = Arc::new(
            GrpcMetaClientPool::new(
                &pd_addrs,
                &config.security,
                config.server.grpc_client.clone(),
                config.meta_service_wait_timeout,
            )
            .await?,
        );
        let metadata_store = Arc::new(MetadataStore::new(&config.storage.metadata_path())?);
        let backend_storage = create_backend_storage(&config.s3, STORAGE_SERVICE_WORKER).await?;
        let frag_batch_uploader = Arc::new(FragBatchUploader::new(
            backend_storage.clone(),
            config.frag_upload.clone(),
        ));
        let shard_writer_manager = DispatchShardWriter::Real(Arc::new(ShardWriterManager::new(
            config.clone(),
            backend_storage,
            meta_client_pool,
            metadata_store,
            frag_path_gen,
            frag_batch_uploader.clone(),
        )));
        Self::new_with_writer(config, shard_writer_manager).await
    }

    pub async fn new_with_writer(
        config: WorkerNodeConfig,
        shard_writer_manager: DispatchShardWriter,
    ) -> Result<Self> {
        info!("worker service created with config={:?}", config);

        config.storage.init_worker_data_dir()?;
        let (meta_client_pool, metadata_store, frag_path_gen, frag_batch_uploader): (
            Arc<dyn MetaClientPool>,
            Arc<MetadataStore>,
            Arc<FragPathGenerator>,
            Arc<FragBatchUploader>,
        ) = match &shard_writer_manager {
            DispatchShardWriter::Real(manager) => (
                manager.meta_client_pool(),
                manager.metadata_store(),
                manager.frag_path_gen(),
                manager.frag_batch_uploader(),
            ),
            DispatchShardWriter::Mock(_) => {
                let pd_addrs = config.server.get_pd_addresses();
                let meta_client_pool: Arc<dyn MetaClientPool> = Arc::new(
                    GrpcMetaClientPool::new(
                        &pd_addrs,
                        &config.security,
                        config.server.grpc_client.clone(),
                        config.meta_service_wait_timeout,
                    )
                    .await?,
                );
                let metadata_store = Arc::new(MetadataStore::new(&config.storage.metadata_path())?);
                let frag_path_gen = Self::create_frag_path_gen(&config).await?;
                let backend_storage =
                    create_backend_storage(&config.s3, STORAGE_SERVICE_WORKER).await?;
                let frag_batch_uploader = Arc::new(FragBatchUploader::new(
                    backend_storage.clone(),
                    config.frag_upload.clone(),
                ));
                (
                    meta_client_pool,
                    metadata_store,
                    frag_path_gen,
                    frag_batch_uploader,
                )
            }
        };

        let worker_node_service = Arc::new(
            WorkerNodeServiceCore::new(
                config.clone(),
                shard_writer_manager,
                config.heartbeat_interval,
                frag_batch_uploader,
                meta_client_pool,
                metadata_store,
                frag_path_gen,
            )
            .await?,
        );

        let status_server = if !config.server.status_addr.is_empty() {
            Some(Self::start_status_server(&config.server)?)
        } else {
            None
        };
        Ok(Self {
            status_server,
            worker_node_service,
            config,
            grpc_handle: None,
        })
    }

    async fn create_frag_path_gen(config: &WorkerNodeConfig) -> Result<Arc<FragPathGenerator>> {
        info!(
            "Starting TSOGenerator, pd_addr={:?}",
            config.server.get_pd_addresses()
        );
        let tso_gen =
            TSOGenerator::new(config.server.get_pd_addresses(), config.security.clone()).await?;

        let frag_base_path = config.storage.frag_base_path();
        info!(
            "Starting FragPathGenerator, frag_data_path={:?}",
            frag_base_path
        );
        Ok(Arc::new(FragPathGenerator::new(
            &frag_base_path,
            Box::new(tso_gen),
        )))
    }

    pub async fn start(&mut self, rx: tokio::sync::oneshot::Receiver<()>) -> Result<()> {
        let addr_str: String = self.config.server.addr.clone();
        let addr = addr_str.parse()?;

        info!(
            "Starting worker service at {}, status_addr={}",
            addr, self.config.server.status_addr
        );

        // Ignore restore error, it will be retried by heartbeat.
        let _ = self
            .worker_node_service
            .restore()
            .await
            .inspect_err(|e| error!("restore failed, err={:?}", e));

        info!("Starting metrics");

        crate::common::server_util::initial_metric();

        self.worker_node_service.start_heartbeat_tasks().await?;

        info!("Starting worker service grpc server at {:?}", addr);
        self.grpc_handle = Some(tokio::spawn(
            Server::builder()
                .add_service(
                    GrpcWorkerServiceServer::new(WorkerGrpcServiceImpl::new(
                        self.worker_node_service.clone(),
                    ))
                    .max_decoding_message_size(
                        self.config.server.grpc.max_decoding_message_size.0 as usize,
                    )
                    .max_encoding_message_size(
                        self.config.server.grpc.max_encoding_message_size.0 as usize,
                    ),
                )
                .serve_with_shutdown(addr, async {
                    let e = rx.await;
                    info!("received stop signal, e={:?}", e);
                    e.unwrap();
                }),
        ));

        Ok(())
    }

    // Take the ownership of `self` and stop all the components.
    pub async fn stop(
        mut self,
        grpc_server_close_tx: tokio::sync::oneshot::Sender<()>,
    ) -> Result<()> {
        info!("WorkerServer is stopping...");
        if let Err(e) = grpc_server_close_tx.send(()) {
            error!("stop grpc error={:?}", e);
        }
        // Stop external services (grpc, http) first
        info!("Grpc server is stopping...");
        if let Some(handle) = self.grpc_handle.take() {
            if let Err(e) = handle.await {
                warn!("Grpc server stopped with err={}", e);
            } else {
                info!("Grpc server stopped successfully");
            }
        }
        if let Some(s) = self.status_server.take() {
            s.stop();
        }
        info!("StatusService stopped successfully");

        info!(
            "Stopping worker service {}",
            self.worker_node_service.config.server.get_advertise_addr()
        );
        self.worker_node_service.shutdown().await?;
        info!(
            "Stopping worker service finished {}",
            self.worker_node_service.config.server.get_advertise_addr()
        );
        Ok(())
    }

    fn start_status_server(config: &ServerConfig) -> Result<StatusServer> {
        let store_path = if config.status_server_store_path.is_empty() {
            tici_common::config::ServerConfig::deduce_status_server_store_path()
        } else {
            config.status_server_store_path.clone().into()
        };

        info!("status_server at path={:?}", store_path);
        let sock_addr = config.status_addr.parse::<SocketAddr>()?;

        let mut status_server = StatusServer::new(store_path.clone(), Operators::None);
        status_server
            .start(sock_addr)
            .with_context(|| format!("Failed to start status server at addr={}", sock_addr))?;
        info!(
            "status_server succeed path={:?}, status_addr={}",
            store_path, config.status_addr
        );
        Ok(status_server)
    }
}