// Copyright 2023 Greptime Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::sync::Arc;

use cache::build_datanode_cache_registry;
use catalog::kvbackend::new_read_only_meta_kv_backend;
use common_base::Plugins;
use common_meta::cache::LayeredCacheRegistryBuilder;
use common_telemetry::info;
use common_version::{short_version, verbose_version};
use datanode::datanode::DatanodeBuilder;
use datanode::service::DatanodeServiceBuilder;
use meta_client::MetaClientType;
use snafu::{OptionExt, ResultExt};
use tracing_appender::non_blocking::WorkerGuard;

use crate::datanode::{APP_NAME, DatanodeOptions, Instance};
use crate::error::{MetaClientInitSnafu, MissingConfigSnafu, Result, StartDatanodeSnafu};
use crate::{create_resource_limit_metrics, log_versions, maybe_activate_heap_profile};

/// Builder for Datanode instance.
pub struct InstanceBuilder {
    guard: Vec<WorkerGuard>,
    opts: DatanodeOptions,
    datanode_builder: DatanodeBuilder,
}

impl InstanceBuilder {
    /// Try to create a new [InstanceBuilder], and do some initialization work like allocating
    /// runtime resources, setting up global logging and plugins, etc.
    pub async fn try_new_with_init(
        mut opts: DatanodeOptions,
        mut plugins: Plugins,
    ) -> Result<Self> {
        let guard = Self::init(&mut opts, &mut plugins).await?;

        let mut datanode_builder = Self::datanode_builder(&opts, &plugins).await?;

        plugins::setup_datanode_plugins_post_build(&mut plugins, &opts.plugins, &datanode_builder)
            .await
            .context(StartDatanodeSnafu)?;
        datanode_builder.set_plugins(plugins);

        Ok(Self {
            guard,
            opts,
            datanode_builder,
        })
    }

    async fn init(opts: &mut DatanodeOptions, plugins: &mut Plugins) -> Result<Vec<WorkerGuard>> {
        let dn_opts = &mut opts.component;
        let guard = common_telemetry::init_global_logging(
            APP_NAME,
            &dn_opts.logging,
            &dn_opts.tracing,
            dn_opts.node_id.map(|x| x.to_string()),
            None,
        );

        common_runtime::init_datanode_runtimes(&opts.runtime);

        crate::options::flush_dropped_plugin_warnings();
        log_versions(verbose_version(), short_version(), APP_NAME);
        maybe_activate_heap_profile(&dn_opts.memory);
        create_resource_limit_metrics(APP_NAME);

        plugins::setup_datanode_plugins_pre_build(plugins, &opts.plugins, dn_opts)
            .await
            .context(StartDatanodeSnafu)?;

        dn_opts.grpc.detect_server_addr();

        info!("Initialized Datanode instance with {:#?}", opts);
        Ok(guard)
    }

    async fn datanode_builder(
        opts: &DatanodeOptions,
        plugins: &Plugins,
    ) -> Result<DatanodeBuilder> {
        let dn_opts = &opts.component;

        let member_id = dn_opts
            .node_id
            .context(MissingConfigSnafu { msg: "'node_id'" })?;
        let meta_client_options = dn_opts.meta_client.as_ref().context(MissingConfigSnafu {
            msg: "meta client options",
        })?;
        let client = meta_client::create_meta_client(
            MetaClientType::Datanode { member_id },
            meta_client_options,
            Some(plugins),
            None,
        )
        .await
        .context(MetaClientInitSnafu)?;

        let backend = new_read_only_meta_kv_backend(client.clone());
        let mut builder = DatanodeBuilder::new(dn_opts.clone(), plugins.clone(), backend.clone());

        let registry = Arc::new(
            LayeredCacheRegistryBuilder::default()
                .add_cache_registry(build_datanode_cache_registry(backend))
                .build(),
        );
        builder
            .with_cache_registry(registry)
            .with_meta_client(client.clone());
        Ok(builder)
    }

    /// Get the mutable builder for Datanode, in case you want to change some fields before the
    /// final construction.
    pub fn mut_datanode_builder(&mut self) -> &mut DatanodeBuilder {
        &mut self.datanode_builder
    }

    /// Try to build the Datanode instance.
    pub async fn build(self) -> Result<Instance> {
        let mut datanode = self
            .datanode_builder
            .build()
            .await
            .context(StartDatanodeSnafu)?;

        let services = DatanodeServiceBuilder::new(&self.opts.component)
            .with_default_grpc_server(&datanode.region_server())
            .enable_http_service()
            .build()
            .context(StartDatanodeSnafu)?;
        datanode.setup_services(services);

        Ok(Instance::new(datanode, self.guard))
    }
}
