// 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::{fmt, mem};

use snafu::ensure;
use tokio::sync::{OwnedSemaphorePermit, TryAcquireError};

use crate::error::{
    MemoryAcquireTimeoutSnafu, MemoryLimitExceededSnafu, MemorySemaphoreClosedSnafu, Result,
};
use crate::manager::{MemoryMetrics, MemoryQuota, UnlimitedMemoryQuota};
use crate::policy::OnExhaustedPolicy;

/// Guard representing a slice of reserved memory.
pub struct MemoryGuard<M: MemoryMetrics> {
    pub(crate) state: GuardState<M>,
}

pub(crate) enum GuardState<M: MemoryMetrics> {
    Released,
    Unlimited {
        quota: UnlimitedMemoryQuota<M>,
        granted_bytes: u64,
    },
    Limited {
        quota: MemoryQuota<M>,
        permit: OwnedSemaphorePermit,
    },
}

impl<M: MemoryMetrics> GuardState<M> {
    fn release(self) {
        match self {
            GuardState::Released => {}
            GuardState::Unlimited {
                quota,
                granted_bytes,
            } => {
                quota.sub_in_use(granted_bytes);
            }
            GuardState::Limited { quota, permit } => {
                quota.release_permit(permit);
            }
        }
    }
}

impl<M: MemoryMetrics> MemoryGuard<M> {
    pub(crate) fn unlimited(quota: UnlimitedMemoryQuota<M>, bytes: u64) -> Self {
        quota.add_in_use(bytes);
        Self {
            state: GuardState::Unlimited {
                quota,
                granted_bytes: bytes,
            },
        }
    }

    pub(crate) fn limited(quota: MemoryQuota<M>, permit: OwnedSemaphorePermit) -> Self {
        Self {
            state: GuardState::Limited { quota, permit },
        }
    }

    /// Returns granted quota in bytes.
    pub fn granted_bytes(&self) -> u64 {
        match &self.state {
            GuardState::Released => 0,
            GuardState::Unlimited { granted_bytes, .. } => *granted_bytes,
            GuardState::Limited { quota, permit } => {
                quota.permits_to_bytes(permit.num_permits() as u32)
            }
        }
    }

    /// Acquires additional memory, waiting if necessary until enough is available.
    ///
    /// On success, merges the new memory into this guard.
    ///
    /// # Errors
    /// - Returns error if requested bytes would exceed the manager's total limit
    /// - Returns error if the semaphore is unexpectedly closed
    pub async fn acquire_additional(&mut self, bytes: u64) -> Result<()> {
        if bytes == 0 {
            return Ok(());
        }

        match &mut self.state {
            GuardState::Released => {
                debug_assert!(false, "released memory guard state should not be reused");
                Ok(())
            }
            GuardState::Unlimited {
                quota,
                granted_bytes,
            } => {
                quota.add_in_use(bytes);
                *granted_bytes = granted_bytes.saturating_add(bytes);
                Ok(())
            }
            GuardState::Limited { quota, permit } => {
                let additional_permits = quota.bytes_to_permits(bytes);
                let current_permits = permit.num_permits() as u32;

                ensure!(
                    current_permits.saturating_add(additional_permits) <= quota.limit_permits,
                    MemoryLimitExceededSnafu {
                        requested_bytes: bytes,
                        limit_bytes: quota.permits_to_bytes(quota.limit_permits)
                    }
                );

                let additional_permit = quota
                    .semaphore
                    .clone()
                    .acquire_many_owned(additional_permits)
                    .await
                    .map_err(|_| MemorySemaphoreClosedSnafu.build())?;

                permit.merge(additional_permit);
                quota.update_in_use_metric();
                Ok(())
            }
        }
    }

    /// Tries to acquire additional memory without waiting.
    ///
    /// On success, merges the new memory into this guard and returns true.
    /// On failure, returns false and leaves this guard unchanged.
    pub fn try_acquire_additional(&mut self, bytes: u64) -> bool {
        if bytes == 0 {
            return true;
        }

        match &mut self.state {
            GuardState::Released => {
                debug_assert!(false, "released memory guard state should not be reused");
                false
            }
            GuardState::Unlimited {
                quota,
                granted_bytes,
            } => {
                quota.add_in_use(bytes);
                *granted_bytes = granted_bytes.saturating_add(bytes);
                true
            }
            GuardState::Limited { quota, permit } => {
                let additional_permits = quota.bytes_to_permits(bytes);

                match quota
                    .semaphore
                    .clone()
                    .try_acquire_many_owned(additional_permits)
                {
                    Ok(additional_permit) => {
                        permit.merge(additional_permit);
                        quota.update_in_use_metric();
                        true
                    }
                    Err(TryAcquireError::NoPermits) | Err(TryAcquireError::Closed) => {
                        quota.metrics.inc_exhausted("try_acquire_additional");
                        false
                    }
                }
            }
        }
    }

    /// Acquires additional memory based on the given policy.
    ///
    /// - For `OnExhaustedPolicy::Wait`: Waits up to the timeout duration for memory to become available
    /// - For `OnExhaustedPolicy::Fail`: Returns immediately if memory is not available
    ///
    /// # Errors
    /// - `MemoryLimitExceeded`: Requested bytes would exceed the total limit (both policies), or memory is currently exhausted (Fail policy only)
    /// - `MemoryAcquireTimeout`: Timeout elapsed while waiting for memory (Wait policy only)
    /// - `MemorySemaphoreClosed`: The internal semaphore is unexpectedly closed (rare, indicates system issue)
    pub async fn acquire_additional_with_policy(
        &mut self,
        bytes: u64,
        policy: OnExhaustedPolicy,
    ) -> Result<()> {
        match policy {
            OnExhaustedPolicy::Wait { timeout } => {
                match tokio::time::timeout(timeout, self.acquire_additional(bytes)).await {
                    Ok(Ok(())) => Ok(()),
                    Ok(Err(e)) => Err(e),
                    Err(_elapsed) => MemoryAcquireTimeoutSnafu {
                        requested_bytes: bytes,
                        waited: timeout,
                    }
                    .fail(),
                }
            }
            OnExhaustedPolicy::Fail => {
                if self.try_acquire_additional(bytes) {
                    Ok(())
                } else {
                    MemoryLimitExceededSnafu {
                        requested_bytes: bytes,
                        limit_bytes: match &self.state {
                            GuardState::Released => 0,
                            GuardState::Unlimited { .. } => 0, // unreachable: unlimited mode always succeeds
                            GuardState::Limited { quota, .. } => {
                                quota.permits_to_bytes(quota.limit_permits)
                            }
                        },
                    }
                    .fail()
                }
            }
        }
    }

    /// Releases a portion of granted memory back to the pool before the guard is dropped.
    ///
    /// Returns true if the release succeeds or is a no-op; false if the request exceeds granted.
    pub fn release_partial(&mut self, bytes: u64) -> bool {
        if bytes == 0 {
            return true;
        }

        match &mut self.state {
            GuardState::Released => true,
            GuardState::Unlimited {
                quota,
                granted_bytes,
            } => {
                if bytes > *granted_bytes {
                    return false;
                }

                quota.sub_in_use(bytes);
                *granted_bytes = granted_bytes.saturating_sub(bytes);
                true
            }
            GuardState::Limited { quota, permit } => {
                let release_permits = quota.bytes_to_permits(bytes);

                match permit.split(release_permits as usize) {
                    Some(released_permit) => {
                        quota.release_permit(released_permit);
                        true
                    }
                    None => false,
                }
            }
        }
    }
}

impl<M: MemoryMetrics> Drop for MemoryGuard<M> {
    fn drop(&mut self) {
        mem::replace(&mut self.state, GuardState::Released).release();
    }
}

impl<M: MemoryMetrics> fmt::Debug for MemoryGuard<M> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("MemoryGuard")
            .field("granted_bytes", &self.granted_bytes())
            .finish()
    }
}
