// 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.

pub mod executor;
mod metadata;

use std::collections::HashMap;

use async_trait::async_trait;
use common_error::ext::BoxedError;
use common_event_recorder::Event;
use common_procedure::error::{ExternalSnafu, FromJsonSnafu, ToJsonSnafu};
use common_procedure::{
    Context as ProcedureContext, Error as ProcedureError, EventContext, EventTrigger, LockKey,
    Procedure, Result as ProcedureResult, Status,
};
use common_telemetry::info;
use common_telemetry::tracing::warn;
#[cfg(feature = "enterprise")]
use common_time::util::current_time_millis;
use common_wal::options::WalOptions;
use serde::{Deserialize, Serialize};
use snafu::{OptionExt, ResultExt};
use store_api::storage::RegionNumber;
use strum::AsRefStr;
use table::metadata::TableId;
use table::table_reference::TableReference;
#[cfg(feature = "enterprise")]
use uuid::Uuid;

use self::executor::DropTableExecutor;
use crate::ddl::DdlContext;
use crate::ddl::event::table::{TableDdlEvent, TableDdlEventType, TableDdlLocator};
use crate::ddl::utils::{convert_region_routes_to_detecting_regions, map_to_procedure_error};
use crate::error::{self, Result};
use crate::key::table_route::TableRouteValue;
use crate::lock_key::{CatalogLock, SchemaLock, TableLock, TableNameLock};
use crate::metrics;
use crate::region_keeper::OperatingRegionGuard;
use crate::rpc::ddl::DropTableTask;
use crate::rpc::router::{RegionRoute, operating_leader_region_roles};

#[cfg(feature = "enterprise")]
fn ensure_retry_later(err: error::Error) -> error::Error {
    if err.is_retry_later() {
        err
    } else {
        error::Error::retry_later(err)
    }
}

pub struct DropTableProcedure {
    /// The context of procedure runtime.
    pub context: DdlContext,
    /// The serializable data.
    pub data: DropTableData,
    /// The guards of opening regions.
    pub(crate) dropping_regions: Vec<OperatingRegionGuard>,
    /// The drop table executor.
    executor: DropTableExecutor,
}

impl DropTableProcedure {
    pub const TYPE_NAME: &'static str = "metasrv-procedure::DropTable";

    pub fn new(task: DropTableTask, context: DdlContext) -> Self {
        let data = DropTableData::new(
            task,
            cfg!(feature = "enterprise") && context.soft_drop_enabled,
            context.soft_drop_retention,
        );
        let executor = data.build_executor();
        Self {
            context,
            data,
            dropping_regions: vec![],
            executor,
        }
    }

    pub fn from_json(json: &str, context: DdlContext) -> ProcedureResult<Self> {
        let data: DropTableData = serde_json::from_str(json).context(FromJsonSnafu)?;
        #[cfg(feature = "enterprise")]
        let mut data = data;
        #[cfg(feature = "enterprise")]
        if data.state == DropTableState::Prepare
            && data.soft_drop_enabled
            && data.dropped_at.is_none()
            && data.soft_drop_retention_millis.is_none()
        {
            data.soft_drop_retention_millis = context
                .soft_drop_retention
                .and_then(|retention| i64::try_from(retention.as_millis()).ok());
        }
        let executor = data.build_executor();

        Ok(Self {
            context,
            data,
            dropping_regions: vec![],
            executor,
        })
    }

    #[cfg(feature = "enterprise")]
    async fn prepare_soft_drop(&mut self) -> Result<()> {
        self.executor
            .check_tombstone_conflict(&self.context, self.data.soft_drop_enabled)
            .await?;
        if self.data.soft_drop_enabled && self.data.dropped_at.is_none() {
            let dropped_at = current_time_millis();
            let retention_millis =
                self.data
                    .soft_drop_retention_millis
                    .context(error::UnexpectedSnafu {
                        err_msg: "Soft-drop retention is missing from the procedure".to_string(),
                    })?;
            let retention_expires_at = dropped_at.checked_add(retention_millis).context(
                error::UnexpectedSnafu {
                    err_msg: format!(
                        "Soft-drop retention deadline overflows i64: dropped_at={dropped_at}, retention_millis={retention_millis}"
                    ),
                },
            )?;
            self.data.dropped_at = Some(dropped_at);
            self.data.retention_expires_at = Some(retention_expires_at);
        }
        if self.data.soft_drop_enabled && self.data.drop_generation.is_none() {
            self.data.drop_generation = Some(Uuid::new_v4().to_string());
        }
        Ok(())
    }

    pub(crate) async fn on_prepare(&mut self) -> Result<Status> {
        if self.executor.on_prepare(&self.context).await?.stop() {
            return Ok(Status::done());
        }
        self.fill_table_metadata().await?;
        #[cfg(feature = "enterprise")]
        self.prepare_soft_drop().await?;
        self.data.state = DropTableState::DeleteMetadata;

        Ok(Status::executing(true))
    }

    /// Register dropping regions if doesn't exist.
    fn register_dropping_regions(&mut self) -> Result<()> {
        let dropping_regions = operating_leader_region_roles(&self.data.physical_region_routes);

        if !self.dropping_regions.is_empty() {
            return Ok(());
        }

        let mut dropping_region_guards = Vec::with_capacity(dropping_regions.len());

        for (region_id, datanode_id, role) in dropping_regions {
            let guard = self
                .context
                .memory_region_keeper
                .register_with_role(datanode_id, region_id, role)
                .context(error::RegionOperatingRaceSnafu {
                    region_id,
                    peer_id: datanode_id,
                })?;
            dropping_region_guards.push(guard);
        }

        self.dropping_regions = dropping_region_guards;
        Ok(())
    }

    #[cfg(not(feature = "enterprise"))]
    async fn delete_metadata(&mut self) -> Result<()> {
        let table_route_value = &TableRouteValue::new(
            self.data.task.table_id,
            // Safety: checked
            self.data.physical_table_id.unwrap(),
            self.data.physical_region_routes.clone(),
        );
        self.executor
            .on_delete_metadata(
                &self.context,
                table_route_value,
                &self.data.region_wal_options,
            )
            .await
    }

    #[cfg(feature = "enterprise")]
    async fn delete_metadata(&mut self) -> Result<()> {
        if !self.data.soft_drop_enabled {
            let table_route_value = &TableRouteValue::new(
                self.data.task.table_id,
                // Safety: checked
                self.data.physical_table_id.unwrap(),
                self.data.physical_region_routes.clone(),
            );
            return self
                .executor
                .on_delete_metadata(
                    &self.context,
                    table_route_value,
                    &self.data.region_wal_options,
                )
                .await;
        }

        let storage = self
            .context
            .table_metadata_manager
            .table_route_manager()
            .table_route_storage();
        storage
            .remap_region_routes(&mut self.data.physical_region_routes)
            .await
            .map_err(ensure_retry_later)?;
        let table_route_value = &TableRouteValue::new(
            self.data.task.table_id,
            // Safety: checked
            self.data.physical_table_id.unwrap(),
            self.data.physical_region_routes.clone(),
        );
        self.executor
            .on_close_regions(
                &self.context.node_manager,
                &self.context.leader_region_registry,
                &self.data.physical_region_routes,
                true,
            )
            .await
            .map_err(ensure_retry_later)?;
        self.executor
            .on_soft_delete_metadata(
                &self.context,
                table_route_value,
                &self.data.region_wal_options,
                self.data.dropped_at,
                self.data.retention_expires_at,
                self.data.drop_generation.as_deref(),
            )
            .await
            .map_err(ensure_retry_later)?;
        self.data.allow_rollback = false;
        Ok(())
    }

    pub(crate) async fn on_delete_metadata(&mut self) -> Result<Status> {
        self.register_dropping_regions()?;
        // TODO(weny): Considers introducing a RegionStatus to indicate the region is dropping.
        let table_id = self.data.table_id();
        self.delete_metadata().await?;
        info!("Deleted table metadata for table {table_id}");
        self.data.state = DropTableState::InvalidateTableCache;
        Ok(Status::executing(true))
    }

    /// Broadcasts invalidate table cache instruction.
    async fn on_broadcast(&mut self) -> Result<Status> {
        let result = self.executor.invalidate_table_cache(&self.context).await;
        #[cfg(feature = "enterprise")]
        if self.data.soft_drop_enabled {
            result.map_err(ensure_retry_later)?;
        } else {
            result?;
        }
        #[cfg(not(feature = "enterprise"))]
        result?;

        self.data.state = DropTableState::DatanodeDropRegions;

        Ok(Status::executing(true))
    }

    pub async fn on_datanode_drop_regions(&mut self, retrying: bool) -> Result<Status> {
        if retrying {
            info!(
                "Remapping region routes addresses for retrying drop regions for table_id: {}",
                self.data.table_id()
            );
            let storage = self
                .context
                .table_metadata_manager
                .table_route_manager()
                .table_route_storage();
            // The peer addresses may change during retries,
            // so we always remap the region routes.
            storage
                .remap_region_routes(&mut self.data.physical_region_routes)
                .await?;
        }

        #[cfg(feature = "enterprise")]
        if self.data.soft_drop_enabled {
            self.context
                .deregister_failure_detectors(convert_region_routes_to_detecting_regions(
                    &self.data.physical_region_routes,
                ))
                .await;
            self.dropping_regions.clear();
            return Ok(Status::done());
        }

        self.executor
            .on_drop_regions(
                &self.context.node_manager,
                &self.context.leader_region_registry,
                &self.data.physical_region_routes,
                false,
                false,
                false,
            )
            .await?;
        self.context
            .deregister_failure_detectors(convert_region_routes_to_detecting_regions(
                &self.data.physical_region_routes,
            ))
            .await;

        self.data.state = DropTableState::DeleteTombstone;
        Ok(Status::executing(true))
    }

    /// Deletes metadata tombstone.
    async fn on_delete_metadata_tombstone(&mut self) -> Result<Status> {
        let table_route_value = &TableRouteValue::new(
            self.data.task.table_id,
            // Safety: checked
            self.data.physical_table_id.unwrap(),
            self.data.physical_region_routes.clone(),
        );
        self.executor
            .on_delete_metadata_tombstone(
                &self.context,
                table_route_value,
                &self.data.region_wal_options,
            )
            .await?;

        self.dropping_regions.clear();
        Ok(Status::done())
    }
}

#[async_trait]
impl Procedure for DropTableProcedure {
    fn type_name(&self) -> &str {
        Self::TYPE_NAME
    }

    fn recover(&mut self) -> ProcedureResult<()> {
        // Only registers regions if the metadata is deleted.
        let register_operating_regions = matches!(
            self.data.state,
            DropTableState::DeleteMetadata
                | DropTableState::InvalidateTableCache
                | DropTableState::DatanodeDropRegions
        );
        if register_operating_regions {
            self.register_dropping_regions()
                .map_err(BoxedError::new)
                .context(ExternalSnafu {
                    clean_poisons: false,
                })?;
        }

        Ok(())
    }

    async fn execute(&mut self, ctx: &ProcedureContext) -> ProcedureResult<Status> {
        let state = &self.data.state;
        let _timer = metrics::METRIC_META_PROCEDURE_DROP_TABLE
            .with_label_values(&[state.as_ref()])
            .start_timer();

        match self.data.state {
            DropTableState::Prepare => self.on_prepare().await,
            DropTableState::DeleteMetadata => self.on_delete_metadata().await,
            DropTableState::InvalidateTableCache => self.on_broadcast().await,
            DropTableState::DatanodeDropRegions => {
                let retrying = ctx.is_retrying().await.unwrap_or(false);
                self.on_datanode_drop_regions(retrying).await
            }
            DropTableState::DeleteTombstone => self.on_delete_metadata_tombstone().await,
        }
        .map_err(map_to_procedure_error)
    }

    fn dump(&self) -> ProcedureResult<String> {
        serde_json::to_string(&self.data).context(ToJsonSnafu)
    }

    fn lock_key(&self) -> LockKey {
        let table_ref = &self.data.table_ref();
        let table_id = self.data.table_id();
        let lock_key = vec![
            CatalogLock::Read(table_ref.catalog).into(),
            SchemaLock::read(table_ref.catalog, table_ref.schema).into(),
            TableNameLock::new(table_ref.catalog, table_ref.schema, table_ref.table).into(),
            TableLock::Write(table_id).into(),
        ];

        LockKey::new(lock_key)
    }

    fn event(&self, ctx: &EventContext<'_>) -> Option<Box<dyn Event>> {
        if !ctx
            .event_type_filter
            .allows(TableDdlEventType::DropTable.as_str())
        {
            return None;
        }
        let task = &self.data.task;
        let locator = TableDdlLocator::new(&task.catalog, &task.schema, &task.table)
            .with_table_id(task.table_id);
        let event = match &ctx.trigger {
            EventTrigger::Submitted => {
                TableDdlEvent::drop_table_submitted(locator, task.drop_if_exists)
            }
            _ => TableDdlEvent::lifecycle(TableDdlEventType::DropTable, [locator]),
        };

        Some(Box::new(event))
    }

    fn rollback_supported(&self) -> bool {
        !matches!(self.data.state, DropTableState::Prepare) && self.data.allow_rollback
    }

    async fn rollback(&mut self, _: &ProcedureContext) -> ProcedureResult<()> {
        warn!(
            "Rolling back the drop table procedure, table: {}",
            self.data.table_id()
        );

        let table_route_value = &TableRouteValue::new(
            self.data.task.table_id,
            // Safety: checked
            self.data.physical_table_id.unwrap(),
            self.data.physical_region_routes.clone(),
        );
        self.executor
            .on_restore_metadata(
                &self.context,
                table_route_value,
                &self.data.region_wal_options,
            )
            .await
            .map_err(ProcedureError::external)?;

        self.dropping_regions.clear();
        Ok(())
    }
}

#[derive(Debug, Serialize, Deserialize)]
pub struct DropTableData {
    pub state: DropTableState,
    pub task: DropTableTask,
    pub physical_region_routes: Vec<RegionRoute>,
    pub physical_table_id: Option<TableId>,
    #[serde(default)]
    pub region_wal_options: HashMap<RegionNumber, WalOptions>,
    #[serde(default)]
    pub allow_rollback: bool,
    #[serde(default)]
    pub soft_drop_enabled: bool,
    #[serde(default)]
    pub dropped_at: Option<i64>,
    #[serde(default)]
    pub retention_expires_at: Option<i64>,
    #[serde(default)]
    pub soft_drop_retention_millis: Option<i64>,
    #[serde(default)]
    pub drop_generation: Option<String>,
}

impl DropTableData {
    pub fn new(
        task: DropTableTask,
        soft_drop_enabled: bool,
        soft_drop_retention: Option<std::time::Duration>,
    ) -> Self {
        Self {
            state: DropTableState::Prepare,
            task,
            physical_region_routes: vec![],
            physical_table_id: None,
            region_wal_options: HashMap::new(),
            allow_rollback: false,
            soft_drop_enabled,
            dropped_at: None,
            retention_expires_at: None,
            soft_drop_retention_millis: soft_drop_retention_millis(soft_drop_retention),
            drop_generation: None,
        }
    }

    fn table_ref(&self) -> TableReference<'_> {
        self.task.table_ref()
    }

    fn table_id(&self) -> TableId {
        self.task.table_id
    }

    fn build_executor(&self) -> DropTableExecutor {
        DropTableExecutor::new(
            self.task.table_name(),
            self.task.table_id,
            self.task.drop_if_exists,
        )
    }
}

#[cfg(feature = "enterprise")]
fn soft_drop_retention_millis(soft_drop_retention: Option<std::time::Duration>) -> Option<i64> {
    soft_drop_retention.and_then(|retention| i64::try_from(retention.as_millis()).ok())
}

#[cfg(not(feature = "enterprise"))]
fn soft_drop_retention_millis(_: Option<std::time::Duration>) -> Option<i64> {
    None
}

/// The state of drop table.
#[derive(Debug, Serialize, Deserialize, AsRefStr, PartialEq)]
pub enum DropTableState {
    /// Prepares to drop the table
    Prepare,
    /// Deletes metadata logically
    DeleteMetadata,
    /// Invalidates Table Cache
    InvalidateTableCache,
    /// Drops regions on Datanode
    DatanodeDropRegions,
    /// Deletes metadata tombstone permanently
    DeleteTombstone,
}
