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

//! Engine event listener for tests.

use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::Duration;

use async_trait::async_trait;
use common_telemetry::info;
use store_api::storage::{FileId, RegionId};
use tokio::sync::{Notify, Semaphore};

use crate::sst::file::RegionFileId;

/// Mito engine background event listener.
#[async_trait]
pub trait EventListener: Send + Sync {
    /// Notifies the listener that a region is flushed successfully.
    fn on_flush_success(&self, region_id: RegionId) {
        let _ = region_id;
    }

    /// Notifies the listener that the engine is stalled.
    fn on_write_stall(&self) {}

    /// Notifies the listener that the region starts to do flush.
    async fn on_flush_begin(&self, region_id: RegionId) {
        let _ = region_id;
    }

    /// Notifies the listener after a flush becomes non-cancellable and before commit.
    async fn on_flush_commit_begin(&self, _region_id: RegionId) {}

    /// Notifies the listener after flush cancellation is requested and its DDL is queued.
    fn on_flush_cancel_requested(&self, _region_id: RegionId) {}

    /// Notifies the listener that the later drop task starts running.
    /// Returns the gc interval if we want to override the default one.
    fn on_later_drop_begin(&self, region_id: RegionId) -> Option<Duration> {
        let _ = region_id;
        None
    }

    /// Notifies the listener that the later drop task of the region is finished.
    fn on_later_drop_end(&self, region_id: RegionId, removed: bool) {
        let _ = region_id;
        let _ = removed;
    }

    /// Notifies the listener that ssts has been merged and the region
    /// is going to update its manifest.
    async fn on_merge_ssts_finished(&self, region_id: RegionId) {
        let _ = region_id;
    }

    /// Notifies the listener that the worker receives requests from the request channel.
    fn on_recv_requests(&self, request_num: usize) {
        let _ = request_num;
    }

    /// Notifies the listener that the file cache is filled when, for example, editing region.
    fn on_file_cache_filled(&self, _file_id: FileId) {}

    /// Notifies the listener that the compaction is scheduled.
    fn on_compaction_scheduled(&self, _region_id: RegionId) {}

    /// Notifies the listener immediately before compaction planning invokes the picker.
    async fn on_compaction_pick_begin(&self, _region_id: RegionId) {}

    /// Notifies the listener after local compaction becomes non-cancellable and before commit.
    async fn on_compaction_commit_begin(&self, _region_id: RegionId) {}

    /// Notifies the listener after compaction results are sent and before pending DDL dispatch.
    async fn on_compaction_result_notified(&self, _region_id: RegionId) {}

    /// Notifies the listener after compaction cancellation is requested and its DDL is queued.
    fn on_compaction_cancel_requested(&self, _region_id: RegionId) {}

    /// Notifies the listener that region starts to send a region change result to worker.
    async fn on_notify_region_change_result_begin(&self, _region_id: RegionId) {}

    /// Notifies the listener that region starts to send a enter staging result to worker.
    async fn on_enter_staging_result_begin(&self, _region_id: RegionId) {}

    /// Notifies the listener that the index build task is executed successfully.
    async fn on_index_build_finish(&self, _region_file_id: RegionFileId) {}

    /// Notifies the listener that the index build task is started.
    async fn on_index_build_begin(&self, _region_file_id: RegionFileId) {}

    /// Notifies the listener that the index build task is aborted.
    async fn on_index_build_abort(&self, _region_file_id: RegionFileId) {}

    /// Notifies the listener after the final source check and before manifest publication.
    async fn on_index_build_before_manifest_commit(&self, _region_file_id: RegionFileId) {}

    /// Notifies the listener after manifest publication and before notifying the worker.
    async fn on_index_build_manifest_committed(&self, _region_file_id: RegionFileId) {}
}

pub type EventListenerRef = Arc<dyn EventListener>;

/// Listener to watch flush events.
#[derive(Default)]
pub struct FlushListener {
    notify: Notify,
    success_count: AtomicUsize,
}

impl FlushListener {
    /// Wait until one flush job is done.
    pub async fn wait(&self) {
        self.notify.notified().await;
    }

    /// Returns the success count.
    pub fn success_count(&self) -> usize {
        self.success_count.load(Ordering::Relaxed)
    }
}

#[async_trait]
impl EventListener for FlushListener {
    fn on_flush_success(&self, region_id: RegionId) {
        info!("Region {} flush successfully", region_id);

        self.success_count.fetch_add(1, Ordering::Relaxed);
        self.notify.notify_one();
    }
}

/// Listener to watch stall events.
#[derive(Default)]
pub struct StallListener {
    notify: Notify,
}

impl StallListener {
    /// Wait for a stall event.
    pub async fn wait(&self) {
        self.notify.notified().await;
    }
}

#[async_trait]
impl EventListener for StallListener {
    fn on_flush_success(&self, _region_id: RegionId) {}

    fn on_write_stall(&self) {
        info!("Engine is stalled");

        self.notify.notify_one();
    }
}

/// Blocks a flush at its begin event until the worker requests flush cancellation.
///
/// Tests use this gate to ensure a flush is active before submitting a lifecycle DDL.
#[derive(Default)]
pub struct FlushCancellationListener {
    flush_started: Notify,
    resume_flush: Notify,
}

impl FlushCancellationListener {
    /// Waits until the flush reaches the cancellation gate.
    pub async fn wait_flush_started(&self) {
        self.flush_started.notified().await;
    }

    /// Allows the blocked flush to continue.
    pub fn resume_flush(&self) {
        self.resume_flush.notify_one();
    }
}

#[async_trait]
impl EventListener for FlushCancellationListener {
    async fn on_flush_begin(&self, region_id: RegionId) {
        info!("Region {} reached the flush cancellation gate", region_id);
        self.flush_started.notify_one();
        self.resume_flush.notified().await;
    }

    fn on_flush_cancel_requested(&self, _region_id: RegionId) {
        self.resume_flush();
    }
}

/// Listener on dropping.
pub struct DropListener {
    gc_duration: Duration,
    notify: Notify,
}

impl DropListener {
    /// Creates a new listener with specific `gc_duration`.
    pub fn new(gc_duration: Duration) -> Self {
        DropListener {
            gc_duration,
            notify: Notify::new(),
        }
    }

    /// Waits until later drop task is done.
    pub async fn wait(&self) {
        self.notify.notified().await;
    }
}

#[async_trait]
impl EventListener for DropListener {
    fn on_later_drop_begin(&self, _region_id: RegionId) -> Option<Duration> {
        Some(self.gc_duration)
    }

    fn on_later_drop_end(&self, _region_id: RegionId, removed: bool) {
        // Asserts result.
        assert!(removed);
        self.notify.notify_one();
    }
}

/// Listener on handling compaction requests.
#[derive(Default)]
pub struct CompactionListener {
    handle_finished_notify: Notify,
    blocker: Notify,
}

impl CompactionListener {
    /// Waits for handling compaction finished request.
    pub async fn wait_handle_finished(&self) {
        self.handle_finished_notify.notified().await;
    }

    /// Wakes up the listener.
    pub fn wake(&self) {
        self.blocker.notify_one();
    }
}

#[async_trait]
impl EventListener for CompactionListener {
    async fn on_merge_ssts_finished(&self, region_id: RegionId) {
        info!("Handle compaction finished request, region {region_id}");

        self.handle_finished_notify.notify_one();

        // Blocks current task.
        self.blocker.notified().await;
    }
}

/// Listener to block on flush and alter.
#[derive(Default)]
pub struct AlterFlushListener {
    flush_begin_notify: Notify,
    block_flush_notify: Notify,
    request_begin_notify: Notify,
    request_count: AtomicUsize,
}

impl AlterFlushListener {
    /// Waits on flush begin.
    pub async fn wait_flush_begin(&self) {
        self.flush_begin_notify.notified().await;
    }

    /// Waits on request begin.
    pub async fn wait_request_begin(&self) {
        self.request_begin_notify.notified().await;
    }

    /// Waits until the worker has received at least `times` request batches.
    pub async fn wait_request_count(&self, times: usize) {
        while self.request_count.load(Ordering::Relaxed) < times {
            self.request_begin_notify.notified().await;
        }
    }

    /// Returns received request batch count.
    pub fn request_count(&self) -> usize {
        self.request_count.load(Ordering::Relaxed)
    }

    /// Continue the flush job.
    pub fn wake_flush(&self) {
        self.block_flush_notify.notify_one();
    }
}

#[async_trait]
impl EventListener for AlterFlushListener {
    async fn on_flush_begin(&self, region_id: RegionId) {
        info!("Wait on notify to start flush for region {}", region_id);

        self.flush_begin_notify.notify_one();
        self.block_flush_notify.notified().await;

        info!("region {} begin flush", region_id);
    }

    fn on_recv_requests(&self, request_num: usize) {
        info!("receive {} request", request_num);

        self.request_count.fetch_add(1, Ordering::Relaxed);
        self.request_begin_notify.notify_one();
    }
}

#[derive(Default)]
pub struct NotifyRegionChangeResultListener {
    notify: Notify,
}

impl NotifyRegionChangeResultListener {
    /// Continue to sending region change result.
    pub fn wake_notify(&self) {
        self.notify.notify_one();
    }
}

#[async_trait]
impl EventListener for NotifyRegionChangeResultListener {
    async fn on_notify_region_change_result_begin(&self, region_id: RegionId) {
        info!(
            "Wait on notify to start notify region change result for region {}",
            region_id
        );
        self.notify.notified().await;
        info!(
            "Continue to sending region change result for region {}",
            region_id
        );
    }
}

#[derive(Default)]
pub struct NotifyEnterStagingResultListener {
    notify: Notify,
}

impl NotifyEnterStagingResultListener {
    /// Continue to sending enter staging result.
    pub fn wake_notify(&self) {
        self.notify.notify_one();
    }
}

#[async_trait]
impl EventListener for NotifyEnterStagingResultListener {
    async fn on_enter_staging_result_begin(&self, region_id: RegionId) {
        info!(
            "Wait on notify to start notify enter staging result for region {}",
            region_id
        );
        self.notify.notified().await;
        info!(
            "Continue to sending enter staging result for region {}",
            region_id
        );
    }
}

#[derive(Default)]
pub struct IndexBuildListener {
    begin_count: AtomicUsize,
    begin_notify: Notify,
    finish_count: AtomicUsize,
    finish_notify: Notify,
    abort_count: AtomicUsize,
    abort_notify: Notify,
    // stop means finished or aborted
    stop_notify: Notify,
}

impl IndexBuildListener {
    /// Wait until index build is done for `times` times.
    pub async fn wait_finish(&self, times: usize) {
        while self.finish_count.load(Ordering::Relaxed) < times {
            self.finish_notify.notified().await;
        }
    }

    /// Wait until index build is stopped for `times` times.
    pub async fn wait_stop(&self, times: usize) {
        while self.finish_count.load(Ordering::Relaxed) + self.abort_count.load(Ordering::Relaxed)
            < times
        {
            self.stop_notify.notified().await;
        }
    }

    /// Wait until index build is begun for `times` times.
    pub async fn wait_begin(&self, times: usize) {
        while self.begin_count.load(Ordering::Relaxed) < times {
            self.begin_notify.notified().await;
        }
    }

    /// Clears the success count.
    pub fn clear_finish_count(&self) {
        self.finish_count.store(0, Ordering::Relaxed);
    }

    /// Returns the success count.
    pub fn finish_count(&self) -> usize {
        self.finish_count.load(Ordering::Relaxed)
    }

    /// Returns the start count.
    pub fn begin_count(&self) -> usize {
        self.begin_count.load(Ordering::Relaxed)
    }

    /// Returns the abort count.
    pub fn abort_count(&self) -> usize {
        self.abort_count.load(Ordering::Relaxed)
    }
}

#[async_trait]
impl EventListener for IndexBuildListener {
    async fn on_index_build_finish(&self, region_file_id: RegionFileId) {
        info!("Region {} index build successfully", region_file_id);
        self.finish_count.fetch_add(1, Ordering::Relaxed);
        self.finish_notify.notify_one();
        self.stop_notify.notify_one();
    }

    async fn on_index_build_begin(&self, region_file_id: RegionFileId) {
        info!("Region {} index build begin", region_file_id);
        self.begin_count.fetch_add(1, Ordering::Relaxed);
        self.begin_notify.notify_one();
    }

    async fn on_index_build_abort(&self, region_file_id: RegionFileId) {
        info!("Region {} index build aborted", region_file_id);
        self.abort_count.fetch_add(1, Ordering::Relaxed);
        self.abort_notify.notify_one();
        self.stop_notify.notify_one();
    }
}

/// A listener that can gate (pause) index build at the begin stage.
///
/// Like [`IndexBuildListener`], it tracks begin/finish/abort counts and exposes
/// wait methods. Additionally, [`GateIndexBuildListener::on_index_build_begin`]
/// blocks until the test calls [`GateIndexBuildListener::release_begin`],
/// enabling deterministic race-condition tests without sleep/flakiness.
///
/// # Usage
///
/// ```ignore
/// let gate = Arc::new(GateIndexBuildListener::default());
/// // pass gate as listener to engine
/// // ... trigger index build ...
/// gate.wait_begin(1).await;       // wait for begin to arrive
/// // ... do something while begin is blocked ...
/// gate.release_begin();           // let begin proceed
/// ```
pub struct GateIndexBuildListener {
    recv_count: AtomicUsize,
    recv_notify: Notify,
    begin_count: AtomicUsize,
    begin_notify: Notify,
    /// Blocks `on_index_build_begin` until released by the test.
    begin_blocker: Semaphore,
    finish_count: AtomicUsize,
    finish_notify: Notify,
    abort_count: AtomicUsize,
    abort_notify: Notify,
    /// stop = finish or abort
    stop_notify: Notify,
}

impl Default for GateIndexBuildListener {
    fn default() -> Self {
        Self {
            recv_count: AtomicUsize::new(0),
            recv_notify: Notify::new(),
            begin_count: AtomicUsize::new(0),
            begin_notify: Notify::new(),
            begin_blocker: Semaphore::new(0),
            finish_count: AtomicUsize::new(0),
            finish_notify: Notify::new(),
            abort_count: AtomicUsize::new(0),
            abort_notify: Notify::new(),
            stop_notify: Notify::new(),
        }
    }
}

impl GateIndexBuildListener {
    /// Wait until the worker has received at least `times` requests.
    pub async fn wait_recv(&self, times: usize) {
        while self.recv_count.load(Ordering::Relaxed) < times {
            self.recv_notify.notified().await;
        }
    }

    /// Wait until index build begin is reached for `times` times.
    pub async fn wait_begin(&self, times: usize) {
        while self.begin_count.load(Ordering::Relaxed) < times {
            self.begin_notify.notified().await;
        }
    }

    /// Wait until index build finish is done for `times` times.
    pub async fn wait_finish(&self, times: usize) {
        while self.finish_count.load(Ordering::Relaxed) < times {
            self.finish_notify.notified().await;
        }
    }

    /// Wait until index build is stopped (finished or aborted) for `times` times.
    pub async fn wait_stop(&self, times: usize) {
        while self.finish_count.load(Ordering::Relaxed) + self.abort_count.load(Ordering::Relaxed)
            < times
        {
            self.stop_notify.notified().await;
        }
    }

    /// Release one blocked `on_index_build_begin` invocation.
    ///
    /// This adds one semaphore permit, so each call releases at most one blocked
    /// begin hook. If a test blocks multiple index build tasks, call this method
    /// once per task. Permits are accumulated if release happens before a task
    /// reaches the gate, but the release order still cannot target a specific
    /// region/file.
    ///
    /// Tests should generally call [`Self::wait_begin`] first and then release
    /// the gate, so the build is known to be blocked before the test performs
    /// the operation under test.
    pub fn release_begin(&self) {
        self.begin_blocker.add_permits(1);
    }

    /// Returns the begin count.
    pub fn begin_count(&self) -> usize {
        self.begin_count.load(Ordering::Relaxed)
    }

    /// Returns the finish count.
    pub fn finish_count(&self) -> usize {
        self.finish_count.load(Ordering::Relaxed)
    }

    /// Returns the abort count.
    pub fn abort_count(&self) -> usize {
        self.abort_count.load(Ordering::Relaxed)
    }

    /// Returns the number of requests received by the worker.
    pub fn recv_count(&self) -> usize {
        self.recv_count.load(Ordering::Relaxed)
    }
}

#[async_trait]
impl EventListener for GateIndexBuildListener {
    fn on_recv_requests(&self, request_num: usize) {
        self.recv_count.fetch_add(request_num, Ordering::Relaxed);
        self.recv_notify.notify_one();
    }

    async fn on_index_build_begin(&self, region_file_id: RegionFileId) {
        info!("Region {} index build begin (gated)", region_file_id);
        self.begin_count.fetch_add(1, Ordering::Relaxed);
        self.begin_notify.notify_one();
        // Block until the test releases the gate.
        let _permit = self
            .begin_blocker
            .acquire()
            .await
            .expect("gate semaphore should not be closed");
    }

    async fn on_index_build_finish(&self, region_file_id: RegionFileId) {
        info!("Region {} index build successfully", region_file_id);
        self.finish_count.fetch_add(1, Ordering::Relaxed);
        self.finish_notify.notify_one();
        self.stop_notify.notify_one();
    }

    async fn on_index_build_abort(&self, region_file_id: RegionFileId) {
        info!("Region {} index build aborted", region_file_id);
        self.abort_count.fetch_add(1, Ordering::Relaxed);
        self.abort_notify.notify_one();
        self.stop_notify.notify_one();
    }
}
