// 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 api::v1::meta::ReconcileRequest;
use async_trait::async_trait;
use catalog::CatalogManagerRef;
use common_base::AffectedRows;
use common_meta::rpc::procedure::{
    GcRegionsRequest as MetaGcRegionsRequest, GcResponse as MetaGcResponse,
    GcTableRequest as MetaGcTableRequest, ManageRegionFollowerRequest, MigrateRegionRequest,
    ProcedureStateResponse,
};
use common_query::Output;
use common_query::error::Result;
use session::context::QueryContextRef;
use store_api::storage::RegionId;
use table::requests::{
    BuildIndexTableRequest, CompactTableRequest, DeleteRequest, FlushTableRequest, InsertRequest,
};
use table::table_name::TableName;

/// A trait for handling table mutations in `QueryEngine`.
#[async_trait]
pub trait TableMutationHandler: Send + Sync {
    /// Inserts rows into the table.
    async fn insert(&self, request: InsertRequest, ctx: QueryContextRef) -> Result<Output>;

    /// Delete rows from the table.
    async fn delete(&self, request: DeleteRequest, ctx: QueryContextRef) -> Result<AffectedRows>;

    /// Trigger a flush task for table.
    async fn flush(&self, request: FlushTableRequest, ctx: QueryContextRef)
    -> Result<AffectedRows>;

    /// Trigger a compaction task for table.
    async fn compact(
        &self,
        request: CompactTableRequest,
        ctx: QueryContextRef,
    ) -> Result<AffectedRows>;

    /// Trigger an index build task for the table.
    async fn build_index(
        &self,
        request: BuildIndexTableRequest,
        ctx: QueryContextRef,
    ) -> Result<AffectedRows>;

    /// Trigger a flush task for a table region.
    async fn flush_region(&self, region_id: RegionId, ctx: QueryContextRef)
    -> Result<AffectedRows>;

    /// Trigger a compaction task for a table region.
    async fn compact_region(
        &self,
        region_id: RegionId,
        ctx: QueryContextRef,
    ) -> Result<AffectedRows>;

    /// Discard all unflushed data from a table region.
    async fn discard_unflushed_data(
        &self,
        region_id: RegionId,
        ctx: QueryContextRef,
    ) -> Result<AffectedRows>;

    /// Discard all unflushed data from all regions of a table.
    async fn discard_unflushed_data_by_table(
        &self,
        table_name: TableName,
        ctx: QueryContextRef,
    ) -> Result<AffectedRows>;
}

/// A trait for handling procedure service requests in `QueryEngine`.
#[async_trait]
pub trait ProcedureServiceHandler: Send + Sync {
    /// Permanently purge a dropped table.
    async fn purge_table(&self, query_ctx: QueryContextRef, table_name: TableName) -> Result<()>;

    /// Migrate a region from source peer to target peer, returns the procedure id if success.
    async fn migrate_region(
        &self,
        query_ctx: QueryContextRef,
        request: MigrateRegionRequest,
    ) -> Result<Option<String>>;

    /// Reconcile a table, database or catalog, returns the procedure id if success.
    async fn reconcile(&self, request: ReconcileRequest) -> Result<Option<String>>;

    /// Query the procedure' state by its id
    async fn query_procedure_state(&self, pid: &str) -> Result<ProcedureStateResponse>;

    /// Manage a region follower to a region.
    async fn manage_region_follower(&self, request: ManageRegionFollowerRequest) -> Result<()>;

    /// Get the catalog manager
    fn catalog_manager(&self) -> &CatalogManagerRef;

    /// Manually trigger GC for specific regions.
    async fn gc_regions(
        &self,
        query_ctx: QueryContextRef,
        request: MetaGcRegionsRequest,
    ) -> Result<MetaGcResponse>;

    /// Manually trigger GC for a table.
    async fn gc_table(
        &self,
        query_ctx: QueryContextRef,
        request: MetaGcTableRequest,
    ) -> Result<MetaGcResponse>;
}

/// This flow service handler is only use for flush flow for now.
#[async_trait]
pub trait FlowServiceHandler: Send + Sync {
    async fn flush(
        &self,
        catalog: &str,
        flow: &str,
        ctx: QueryContextRef,
    ) -> Result<api::v1::flow::FlowResponse>;
}

pub type TableMutationHandlerRef = Arc<dyn TableMutationHandler>;

pub type ProcedureServiceHandlerRef = Arc<dyn ProcedureServiceHandler>;

pub type FlowServiceHandlerRef = Arc<dyn FlowServiceHandler>;
