use alloc::boxed::Box;

use alloc::vec::Vec;

use core::mem;

#[cfg(feature = "std")]

use std::sync::{RwLock, RwLockReadGuard};

use pki_types::UnixTime;

use crate::lock::{Mutex, MutexGuard};

use crate::server::ProducesTickets;

#[cfg(not(feature = "std"))]

use crate::time_provider::TimeProvider;

use crate::{rand, Error};

#[derive(Debug)]

pub(crate) struct TicketSwitcherState {

    next: Option<Box<dyn ProducesTickets>>,

    current: Box<dyn ProducesTickets>,

    previous: Option<Box<dyn ProducesTickets>>,

    next_switch_time: u64,

}

/// A ticketer that has a 'current' sub-ticketer and a single

/// 'previous' ticketer.  It creates a new ticketer every so

/// often, demoting the current ticketer.

#[cfg_attr(feature = "std", derive(Debug))]

pub struct TicketSwitcher {

    pub(crate) generator: fn() -> Result<Box<dyn ProducesTickets>, rand::GetRandomFailed>,

    lifetime: u32,

    state: Mutex<TicketSwitcherState>,

    #[cfg(not(feature = "std"))]

    time_provider: &'static dyn TimeProvider,

}

impl TicketSwitcher {

    /// Creates a new `TicketSwitcher`, which rotates through sub-ticketers

    /// based on the passage of time.

    ///

    /// `lifetime` is in seconds, and is how long the current ticketer

    /// is used to generate new tickets.  Tickets are accepted for no

    /// longer than twice this duration.  `generator` produces a new

    /// `ProducesTickets` implementation.

    #[cfg(feature = "std")]

    #[deprecated(note = "use TicketRotator instead")]

    pub fn new(

        lifetime: u32,

        generator: fn() -> Result<Box<dyn ProducesTickets>, rand::GetRandomFailed>,

    ) -> Result<Self, Error> {

        Ok(Self {

            generator,

            lifetime,

            state: Mutex::new(TicketSwitcherState {

                next: Some(generator()?),

                current: generator()?,

                previous: None,

                next_switch_time: UnixTime::now()

                    .as_secs()

                    .saturating_add(u64::from(lifetime)),

            }),

        })

    }

    /// Creates a new `TicketSwitcher`, which rotates through sub-ticketers

    /// based on the passage of time.

    ///

    /// `lifetime` is in seconds, and is how long the current ticketer

    /// is used to generate new tickets.  Tickets are accepted for no

    /// longer than twice this duration.  `generator` produces a new

    /// `ProducesTickets` implementation.

    #[cfg(not(feature = "std"))]

    pub fn new<M: crate::lock::MakeMutex>(

        lifetime: u32,

        generator: fn() -> Result<Box<dyn ProducesTickets>, rand::GetRandomFailed>,

        time_provider: &'static dyn TimeProvider,

    ) -> Result<Self, Error> {

        Ok(Self {

            generator,

            lifetime,

            state: Mutex::new::<M>(TicketSwitcherState {

                next: Some(generator()?),

                current: generator()?,

                previous: None,

                next_switch_time: time_provider

                    .current_time()

                    .unwrap()

                    .as_secs()

                    .saturating_add(u64::from(lifetime)),

            }),

            time_provider,

        })

    }

    /// If it's time, demote the `current` ticketer to `previous` (so it

    /// does no new encryptions but can do decryption) and use next for a

    /// new `current` ticketer.

    ///

    /// Calling this regularly will ensure timely key erasure.  Otherwise,

    /// key erasure will be delayed until the next encrypt/decrypt call.

    ///

    /// For efficiency, this is also responsible for locking the state mutex

    /// and returning the mutexguard.

    pub(crate) fn maybe_roll(&self, now: UnixTime) -> Option<MutexGuard<'_, TicketSwitcherState>> {

        // The code below aims to make switching as efficient as possible

        // in the common case that the generator never fails. To achieve this

        // we run the following steps:

        //  1. If no switch is necessary, just return the mutexguard

        //  2. Shift over all of the ticketers (so current becomes previous,

        //     and next becomes current). After this, other threads can

        //     start using the new current ticketer.

        //  3. unlock mutex and generate new ticketer.

        //  4. Place new ticketer in next and return current

        //

        // There are a few things to note here. First, we don't check whether

        // a new switch might be needed in step 4, even though, due to locking

        // and entropy collection, significant amounts of time may have passed.

        // This is to guarantee that the thread doing the switch will eventually

        // make progress.

        //

        // Second, because next may be None, step 2 can fail. In that case

        // we enter a recovery mode where we generate 2 new ticketers, one for

        // next and one for the current ticketer. We then take the mutex a

        // second time and redo the time check to see if a switch is still

        // necessary.

        //

        // This somewhat convoluted approach ensures good availability of the

        // mutex, by ensuring that the state is usable and the mutex not held

        // during generation. It also ensures that, so long as the inner

        // ticketer never generates panics during encryption/decryption,

        // we are guaranteed to never panic when holding the mutex.

        let now = now.as_secs();

        let mut are_recovering = false; // Are we recovering from previous failure?

        {

            // Scope the mutex so we only take it for as long as needed

            let mut state = self.state.lock()?;

            // Fast path in case we do not need to switch to the next ticketer yet

            if now <= state.next_switch_time {

                return Some(state);

            }

            // Make the switch, or mark for recovery if not possible

            if let Some(next) = state.next.take() {

                state.previous = Some(mem::replace(&mut state.current, next));

                state.next_switch_time = now.saturating_add(u64::from(self.lifetime));

            } else {

                are_recovering = true;

            }

        }

        // We always need a next, so generate it now

        let next = (self.generator)().ok()?;

        if !are_recovering {

            // Normal path, generate new next and place it in the state

            let mut state = self.state.lock()?;

            state.next = Some(next);

            Some(state)

        } else {

            // Recovering, generate also a new current ticketer, and modify state

            // as needed. (we need to redo the time check, otherwise this might

            // result in very rapid switching of ticketers)

            let new_current = (self.generator)().ok()?;

            let mut state = self.state.lock()?;

            state.next = Some(next);

            if now > state.next_switch_time {

                state.previous = Some(mem::replace(&mut state.current, new_current));

                state.next_switch_time = now.saturating_add(u64::from(self.lifetime));

            }

            Some(state)

        }

    }

}

impl ProducesTickets for TicketSwitcher {

    fn lifetime(&self) -> u32 {

        self.lifetime * 2

    }

    fn enabled(&self) -> bool {

        true

    }

    fn encrypt(&self, message: &[u8]) -> Option<Vec<u8>> {

        #[cfg(feature = "std")]

        let now = UnixTime::now();

        #[cfg(not(feature = "std"))]

        let now = self

            .time_provider

            .current_time()

            .unwrap();

        self.maybe_roll(now)?

            .current

            .encrypt(message)

    }

    fn decrypt(&self, ciphertext: &[u8]) -> Option<Vec<u8>> {

        #[cfg(feature = "std")]

        let now = UnixTime::now();

        #[cfg(not(feature = "std"))]

        let now = self

            .time_provider

            .current_time()

            .unwrap();

        let state = self.maybe_roll(now)?;

        // Decrypt with the current key; if that fails, try with the previous.

        state

            .current

            .decrypt(ciphertext)

            .or_else(|| {

                state

                    .previous

                    .as_ref()

                    .and_then(|previous| previous.decrypt(ciphertext))

            })

    }

}

#[cfg(not(feature = "std"))]

impl core::fmt::Debug for TicketSwitcher {

    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {

        f.debug_struct("TicketSwitcher")

            .field("generator", &self.generator)

            .field("lifetime", &self.lifetime)

            .field("state", &**self.state.lock().unwrap())

            .finish()

    }

}

#[cfg(feature = "std")]

#[derive(Debug)]

pub(crate) struct TicketRotatorState {

    current: Box<dyn ProducesTickets>,

    previous: Option<Box<dyn ProducesTickets>>,

    next_switch_time: u64,

}

/// A ticketer that has a 'current' sub-ticketer and a single

/// 'previous' ticketer.  It creates a new ticketer every so

/// often, demoting the current ticketer.

#[cfg(feature = "std")]

pub struct TicketRotator {

    pub(crate) generator: fn() -> Result<Box<dyn ProducesTickets>, rand::GetRandomFailed>,

    lifetime: u32,

    state: RwLock<TicketRotatorState>,

}

#[cfg(feature = "std")]

impl TicketRotator {

    /// Creates a new `TicketRotator`, which rotates through sub-ticketers

    /// based on the passage of time.

    ///

    /// `lifetime` is in seconds, and is how long the current ticketer

    /// is used to generate new tickets.  Tickets are accepted for no

    /// longer than twice this duration.  `generator` produces a new

    /// `ProducesTickets` implementation.

    pub fn new(

        lifetime: u32,

        generator: fn() -> Result<Box<dyn ProducesTickets>, rand::GetRandomFailed>,

    ) -> Result<Self, Error> {

        Ok(Self {

            generator,

            lifetime,

            state: RwLock::new(TicketRotatorState {

                current: generator()?,

                previous: None,

                next_switch_time: UnixTime::now()

                    .as_secs()

                    .saturating_add(u64::from(lifetime)),

            }),

        })

    }

    /// If it's time, demote the `current` ticketer to `previous` (so it

    /// does no new encryptions but can do decryption) and replace it

    /// with a new one.

    ///

    /// Calling this regularly will ensure timely key erasure.  Otherwise,

    /// key erasure will be delayed until the next encrypt/decrypt call.

    ///

    /// For efficiency, this is also responsible for locking the state rwlock

    /// and returning it for read.

    pub(crate) fn maybe_roll(

        &self,

        now: UnixTime,

    ) -> Option<RwLockReadGuard<'_, TicketRotatorState>> {

        let now = now.as_secs();

        // Fast, common, & read-only path in case we do not need to switch

        // to the next ticketer yet

        {

            let read = self.state.read().ok()?;

            if now <= read.next_switch_time {

                return Some(read);

            }

        }

        // We need to switch ticketers, and make a new one.

        // Generate a potential "next" ticketer outside the lock.

        let next = (self.generator)().ok()?;

        let mut write = self.state.write().ok()?;

        if now <= write.next_switch_time {

            // Another thread beat us to it.  Nothing to do.

            drop(write);

            return self.state.read().ok();

        }

        // Now we have:

        // - confirmed we need rotation

        // - confirmed we are the thread that will do it

        // - successfully made the replacement ticketer

        write.previous = Some(mem::replace(&mut write.current, next));

        write.next_switch_time = now.saturating_add(u64::from(self.lifetime));

        drop(write);

        self.state.read().ok()

    }

}

#[cfg(feature = "std")]

impl ProducesTickets for TicketRotator {

    fn lifetime(&self) -> u32 {

        self.lifetime * 2

    }

    fn enabled(&self) -> bool {

        true

    }

    fn encrypt(&self, message: &[u8]) -> Option<Vec<u8>> {

        self.maybe_roll(UnixTime::now())?

            .current

            .encrypt(message)

    }

    fn decrypt(&self, ciphertext: &[u8]) -> Option<Vec<u8>> {

        let state = self.maybe_roll(UnixTime::now())?;

        // Decrypt with the current key; if that fails, try with the previous.

        state

            .current

            .decrypt(ciphertext)

            .or_else(|| {

                state

                    .previous

                    .as_ref()

                    .and_then(|previous| previous.decrypt(ciphertext))

            })

    }

}

#[cfg(feature = "std")]

impl core::fmt::Debug for TicketRotator {

    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {

        f.debug_struct("TicketRotator")

            .finish_non_exhaustive()

    }

}