#![doc = include_str!("../README.md")]
#![deny(missing_docs)]
#![cfg_attr(not(any(feature = "std", test)), no_std)]

extern crate alloc;

#[cfg(test)]
mod mock;

#[expect(clippy::cast_possible_truncation)]
#[frame_support::pallet]
mod pallet {
  use frame_system::pallet_prelude::*;
  use frame_support::pallet_prelude::*;

  use serai_abi::{
    primitives::{
      prelude::*,
      dex::{Error as PrimitivesError, Reserves, Premise},
    },
    Event,
  };

  use serai_core_pallet::Pallet as Core;
  type Coins<T> = serai_coins_pallet::Pallet<T, serai_coins_pallet::CoinsInstance>;
  type LiquidityTokens<T> =
    serai_coins_pallet::Pallet<T, serai_coins_pallet::LiquidityTokensInstance>;

  use super::*;

  /// The configuration of this pallet.
  #[pallet::config]
  pub trait Config:
    frame_system::Config
    + serai_core_pallet::Config
    + serai_coins_pallet::Config<serai_coins_pallet::CoinsInstance>
    + serai_coins_pallet::Config<serai_coins_pallet::LiquidityTokensInstance>
  {
  }

  /// An error incurred.
  #[pallet::error]
  pub enum Error<T> {
    /// The effected/would-be-used liquidity is invalid.
    InvalidLiquidity,
    /// An arithmetic overflow occurred.
    Overflow,
    /// An arithmetic underflow occured.
    Underflow,
    /// The requested swap wasn't satisfied.
    Unsatisfied,
    /// The swap was from the coin it's to.
    FromToSelf,
  }

  impl<T> From<PrimitivesError> for Error<T> {
    fn from(error: PrimitivesError) -> Error<T> {
      match error {
        PrimitivesError::Overflow => Error::Overflow,
        PrimitivesError::Underflow => Error::Underflow,
        PrimitivesError::KInvariant => Error::Unsatisfied,
      }
    }
  }

  #[pallet::storage]
  type LpFeeInThousandths<T: Config> = StorageMap<_, Identity, ExternalCoin, u8, OptionQuery>;

  /// The Pallet struct.
  #[pallet::pallet]
  pub struct Pallet<T>(_);

  impl<T: Config> Pallet<T> {
    fn emit_event(event: Event) {
      Core::<T>::emit_event(event)
    }
  }

  const MINIMUM_LIQUIDITY: u64 = 1 << 16;

  #[pallet::call]
  impl<T: Config> Pallet<T> {
    /// Add liquidity.
    #[pallet::call_index(0)]
    #[pallet::weight((0, DispatchClass::Normal))] // TODO
    pub fn add_liquidity(
      origin: OriginFor<T>,
      external_coin: ExternalCoin,
      sri_intended: Amount,
      external_coin_intended: Amount,
      sri_minimum: Amount,
      external_coin_minimum: Amount,
    ) -> DispatchResult {
      let from = ensure_signed(origin)?;

      let pool = serai_abi::dex::address(external_coin);
      let supply = LiquidityTokens::<T>::supply(Coin::from(external_coin)).0;

      let (sri_actual, external_coin_actual, liquidity) = if supply == 0 {
        let sri_actual = sri_intended;
        let external_coin_actual = external_coin_intended;
        let liquidity = Amount(
          u64::try_from((u128::from(sri_actual.0) * u128::from(external_coin_actual.0)).isqrt())
            .map_err(|_| Error::<T>::Overflow)?,
        );
        if liquidity.0 < MINIMUM_LIQUIDITY {
          Err(Error::<T>::InvalidLiquidity)?;
        }
        (sri_intended, external_coin_intended, liquidity)
      } else {
        let reserves = Reserves {
          sri: Coins::<T>::balance(pool, Coin::Serai),
          external_coin: Coins::<T>::balance(pool, Coin::from(external_coin)),
        };

        let (sri_actual, external_coin_actual) = {
          let (sri_optimal, external_coin_optimal) = (
            Premise::establish(Coin::from(external_coin), Coin::Serai)
              .expect("ext, sri satisfies sri ^ sri")
              .quote_for_in(reserves, external_coin_intended)
              .map_err(Error::<T>::from)?,
            Premise::establish(Coin::Serai, Coin::from(external_coin))
              .expect("sri, ext satisfies sri ^ sri")
              .quote_for_in(reserves, sri_intended)
              .map_err(Error::<T>::from)?,
          );
          if sri_optimal < sri_intended {
            if sri_optimal < sri_minimum {
              Err(Error::<T>::Unsatisfied)?;
            }
            (sri_optimal, external_coin_intended)
          } else {
            if external_coin_optimal < external_coin_minimum {
              Err(Error::<T>::Unsatisfied)?;
            }
            (sri_intended, external_coin_optimal)
          }
        };

        let liquidity = {
          let supply = u128::from(supply);
          let sri_liquidity =
            u64::try_from((u128::from(sri_actual.0) * supply) / u128::from(reserves.sri.0))
              .map_err(|_| Error::<T>::Overflow)?;
          let external_coin_liquidity = u64::try_from(
            (u128::from(external_coin_actual.0) * supply) / u128::from(reserves.external_coin.0),
          )
          .map_err(|_| Error::<T>::Overflow)?;
          Amount(sri_liquidity.min(external_coin_liquidity))
        };

        (sri_actual, external_coin_actual, liquidity)
      };

      Coins::<T>::transfer_fn(
        from,
        pool.into(),
        Balance { coin: Coin::Serai, amount: sri_actual },
      )?;
      Coins::<T>::transfer_fn(
        from,
        pool.into(),
        Balance { coin: Coin::from(external_coin), amount: external_coin_actual },
      )?;
      LiquidityTokens::<T>::mint(
        from,
        Balance { coin: Coin::from(external_coin), amount: liquidity },
      )?;

      // TODO: Event

      Ok(())
    }

    /// Transfer these liquidity tokens to the specified address.
    #[pallet::call_index(1)]
    #[pallet::weight((0, DispatchClass::Normal))] // TODO
    pub fn transfer_liquidity(
      origin: OriginFor<T>,
      to: SeraiAddress,
      liquidity_tokens: ExternalBalance,
    ) -> DispatchResult {
      let from = ensure_signed(origin)?;
      LiquidityTokens::<T>::transfer_fn(from, to.into(), liquidity_tokens.into())?;

      // TODO: Event

      Ok(())
    }

    /// Remove liquidity.
    #[pallet::call_index(2)]
    #[pallet::weight((0, DispatchClass::Normal))] // TODO
    pub fn remove_liquidity(
      origin: OriginFor<T>,
      liquidity_tokens: ExternalBalance,
      sri_minimum: Amount,
      external_coin_minimum: Amount,
    ) -> DispatchResult {
      let from = ensure_signed(origin)?;

      let external_coin = liquidity_tokens.coin;
      let pool = serai_abi::dex::address(external_coin);
      let supply = LiquidityTokens::<T>::supply(Coin::from(external_coin)).0;
      if supply.saturating_sub(liquidity_tokens.amount.0) < MINIMUM_LIQUIDITY {
        Err(Error::<T>::InvalidLiquidity)?;
      }
      let supply = u128::from(supply);

      let reserves = Reserves {
        sri: Coins::<T>::balance(pool, Coin::Serai),
        external_coin: Coins::<T>::balance(pool, Coin::from(external_coin)),
      };
      let sri_amount =
        (u128::from(liquidity_tokens.amount.0) * u128::from(reserves.sri.0)) / supply;
      let sri_amount = Amount(u64::try_from(sri_amount).map_err(|_| Error::<T>::Overflow)?);
      let external_coin_amount =
        (u128::from(liquidity_tokens.amount.0) * u128::from(reserves.external_coin.0)) / supply;
      let external_coin_amount =
        Amount(u64::try_from(external_coin_amount).map_err(|_| Error::<T>::Overflow)?);
      if (sri_amount < sri_minimum) || (external_coin_amount < external_coin_minimum) {
        Err(Error::<T>::Unsatisfied)?;
      }

      LiquidityTokens::<T>::burn_fn(from, liquidity_tokens.into())?;
      Coins::<T>::transfer_fn(
        from,
        pool.into(),
        Balance { coin: Coin::Serai, amount: sri_amount },
      )?;
      Coins::<T>::transfer_fn(
        from,
        pool.into(),
        Balance { coin: Coin::from(external_coin), amount: external_coin_amount },
      )?;

      // TODO: Event

      Ok(())
    }

    /// Swap an exact amount of coins.
    #[pallet::call_index(3)]
    #[pallet::weight((0, DispatchClass::Normal))] // TODO
    pub fn swap(
      origin: OriginFor<T>,
      coins_to_swap: Balance,
      minimum_to_receive: Balance,
    ) -> DispatchResult {
      let origin = SeraiAddress::from(ensure_signed(origin)?);

      let mut transfer_from = origin;
      let mut next_amount = coins_to_swap.amount;

      let swaps = Premise::route(coins_to_swap.coin, minimum_to_receive.coin)
        .ok_or(Error::<T>::FromToSelf)?;
      for swap in &swaps {
        let external_coin = swap.external_coin();
        let pool = serai_abi::dex::address(external_coin);

        // Fetch the pool's reserves
        let reserves = Reserves {
          sri: Coins::<T>::balance(pool, Coin::Serai),
          external_coin: Coins::<T>::balance(pool, Coin::from(external_coin)),
        };
        if (reserves.sri == Amount(0)) || (reserves.external_coin == Amount(0)) {
          Err(Error::<T>::InvalidLiquidity)?;
        }

        // Transfer from the prior (the originating account or pool) to the current pool
        /*
          This is impossible to reach, yet would cause the amount _yet to be transferred out_ to
          be credited both as part of the reserves _and_ the amount in if violated.
        */
        assert!(transfer_from != pool, "swap routed from a coin to itself");
        Coins::<T>::transfer_fn(
          transfer_from.into(),
          pool.into(),
          Balance { coin: swap.r#in(), amount: next_amount },
        )?;

        // Update the current status
        transfer_from = pool;
        next_amount = swap.quote_for_in(reserves, next_amount).map_err(Error::<T>::from)?;
      }

      // Check the amount meets the expectation
      if next_amount.0 < minimum_to_receive.amount.0 {
        Err(Error::<T>::Unsatisfied)?;
      }

      // Transfer the resulting coins to the origin
      Coins::<T>::transfer_fn(
        transfer_from.into(),
        origin.into(),
        Balance { coin: minimum_to_receive.coin, amount: next_amount },
      )?;

      // TODO: Event

      Ok(())
    }

    /// Swap for an exact amount of coins.
    #[pallet::call_index(4)]
    #[pallet::weight((0, DispatchClass::Normal))] // TODO
    pub fn swap_for(
      origin: OriginFor<T>,
      coins_to_receive: Balance,
      maximum_to_swap: Balance,
    ) -> DispatchResult {
      let origin = SeraiAddress::from(ensure_signed(origin)?);

      let mut transfer_to = origin;
      let mut next_amount = coins_to_receive.amount;

      let swaps = Premise::route(maximum_to_swap.coin, coins_to_receive.coin)
        .ok_or(Error::<T>::FromToSelf)?;
      let mut i = swaps.len();
      while {
        i -= 1;
        let swap = swaps[i];

        let external_coin = swap.external_coin();
        let pool = serai_abi::dex::address(external_coin);

        // Fetch the pool's reserves
        let reserves = Reserves {
          sri: Coins::<T>::balance(pool, Coin::Serai),
          external_coin: Coins::<T>::balance(pool, Coin::from(external_coin)),
        };

        /*
          Transfer the output requested.

          While this violates the traditional Checks-Effects-Interactions pattern, in a way
          favoring the caller, this function is within a `transactional` layer
          (due to being a call) making all its DB operations only persisted upon success.
        */
        /*
          This is impossible to reach, yet ensures the amount not yet transferred in isn't
          excluded when determining the reserves on the next iteration.
        */
        assert!(transfer_to != pool, "swap routed to a coin from itself");
        Coins::<T>::transfer_fn(
          pool.into(),
          transfer_to.into(),
          Balance { coin: swap.out(), amount: next_amount },
        )?;

        transfer_to = pool;
        next_amount = swap.quote_for_out(reserves, next_amount).map_err(Error::<T>::from)?;

        i != 0
      } {}

      // Check the amount meets the expectation
      if next_amount.0 > maximum_to_swap.amount.0 {
        Err(Error::<T>::Unsatisfied)?;
      }

      // Transfer the necessary coins from the origin
      Coins::<T>::transfer_fn(
        origin.into(),
        transfer_to.into(),
        Balance { coin: maximum_to_swap.coin, amount: next_amount },
      )?;

      // TODO: Event

      Ok(())
    }
  }
}

pub use pallet::*;