{-# LANGUAGE TemplateHaskell #-}

module UTxO.Generator (
    -- * Inputs
    -- ** State
    GenInpState(..)
  , gisUtxo
  , initInpState
    -- ** Generator
  , RemoveUsedInputs(..)
  , GenInput
  , genInput
    -- * Outputs
    -- ** Parameters
  , GenOutParams(..)
  , defOutParams
    -- ** State
  , GenOutState(..)
  , gosAvailable
  , initOutState
    -- ** Generator
  , GenOutput
  , genOutput
    -- * Transactions
    -- ** Parameters
  , GenTrParams(..)
  , defTrParams
    -- ** State
  , GenTrState(..)
  , gtsInpState
  , gtsNextHash
  , initTrState
    -- ** Generator
  , GenTransaction
  , MakeOutputsAvailable(..)
  , genTransaction
    -- * Chains
    -- ** Params
  , GenChainParams(..)
  , defChainParams
    -- ** Generator
  , genChain
    -- * Auxiliary
  , replicateAtMostM
  ) where

import           Universum

import           Control.Lens (zoom, (%=), (.=), (<<+=))
import           Control.Lens.TH (makeLenses)
import qualified Data.Set as Set
import           Pos.Core (maxCoinVal)
import           Pos.Core.Chrono
import           Test.QuickCheck

import           UTxO.DSL

{-------------------------------------------------------------------------------
  Generate transaction inputs
-------------------------------------------------------------------------------}

-- | State needed for input generation
data GenInpState h a = GenInpState {
      -- | Available UTxO
      _gisUtxo :: Utxo h a
    }

makeLenses ''GenInpState

-- | Initial 'GenInpState'
initInpState :: Utxo h a -> GenInpState h a
initInpState utxo = GenInpState {
      _gisUtxo = utxo
    }

-- | Input generator
type GenInput h a = StateT (GenInpState h a) Gen

-- | Should we remove used inputs?
data RemoveUsedInputs = RemoveUsedInputs | DontRemoveUsedInputs

-- | Try to generate an input
--
-- Returns nothing if the utxo is empty.
genInput :: Hash h a
         => RemoveUsedInputs
         -> Set (Input h a) -- Inputs to avoid
         -> GenInput h a (Maybe (Input h a, Output h a))
genInput removeUsedInputs notThese = do
    utxo <- utxoRemoveInputs notThese <$> use gisUtxo
    if utxoNull utxo
      then return Nothing
      else do
        (inp, out) <- lift $ elements (utxoToList utxo)

        case removeUsedInputs of
          DontRemoveUsedInputs -> return ()
          RemoveUsedInputs     -> gisUtxo .= utxoRemoveInputs (Set.singleton inp) utxo

        return $ Just (inp, out)

-- | Generate up to @n@ distinct inputs
genDistinctInputs :: forall h a. Hash h a
                  => RemoveUsedInputs
                  -> Set (Input h a) -- Inputs to avoid
                  -> Int
                  -> GenInput h a [(Input h a, Output h a)]
genDistinctInputs removeUsedInputs = go
  where
    go :: Set (Input h a) -> Int -> GenInput h a [(Input h a, Output h a)]
    go _        0 = return []
    go notThese n = do
        mInp <- genInput removeUsedInputs notThese
        case mInp of
          Nothing ->
            return []
          Just (inp, out) ->
            -- Removing used inputs or not, don't select same input twice
            ((inp, out) :) <$> go (Set.insert inp notThese) (n - 1)

{-------------------------------------------------------------------------------
  Generate transaction outputs
-------------------------------------------------------------------------------}

-- | Parameters for output generation
data GenOutParams a = GenOutParams {
      -- | Addresses we can generate outputs to
      gopAddresses :: [a]
    }

-- | Default 'GenOutParams'
defOutParams :: [a] -- ^ Addresses we can generate outputs to
             -> GenOutParams a
defOutParams addresses = GenOutParams {
      gopAddresses = addresses
    }

-- | State needed for output generation
data GenOutState = GenOutState {
      -- | Value left
      _gosAvailable :: Value
    }

makeLenses ''GenOutState

-- | Initial 'GenOutState'
initOutState :: Value -> GenOutState
initOutState available = GenOutState {
      _gosAvailable = available
    }

-- | Output generator
type GenOutput = StateT GenOutState Gen

-- | Try to generate transaction output
--
-- Returns nothing if there is no balance left.
genOutput :: GenOutParams a -> GenOutput (Maybe (Output h a))
genOutput GenOutParams{..} = do
    available <- use gosAvailable
    if available == 0
      then return Nothing
      else do
        addr <- lift $ elements gopAddresses
        val  <- lift $ choose (1, min available maxCoinVal)
        gosAvailable .= available - val
        return $ Just (Output addr val)

{-------------------------------------------------------------------------------
  Generate transaction
-------------------------------------------------------------------------------}

-- | Parameters for transaction generation
data GenTrParams a = GenTrParams {
      -- | Maximum number of inputs
      --
      -- Generation is biased towards smaller values
      gtpMaxNumInputs  :: Int

      -- | Maximum number of outputs
      --
      -- Generation does a uniform draw.
    , gtpMaxNumOutputs :: Int

      -- | Fee model
      --
      -- Provide fee given number of inputs and outputs
    , gtpEstimateFee   :: Int -> Int -> Value

      -- | Output parameters
    , gtpOutParams     :: GenOutParams a
    }

-- | Default 'GenTrParams'
defTrParams :: (Int -> Int -> Value)    -- ^ Fee model
            -> [a]                      -- ^ Addresses we can generate outputs to
            -> GenTrParams a
defTrParams feeModel addresses = GenTrParams {
      gtpMaxNumInputs  = 3
    , gtpMaxNumOutputs = 3
    , gtpEstimateFee   = feeModel
    , gtpOutParams     = defOutParams addresses
    }

-- | Should the outputs of the transaction be made available in the
-- input generation state?
data MakeOutputsAvailable =
    -- | Yes, make outputs available
    --
    -- This means that the next generation can refer to outputs of the
    -- previous generated transaction.
    MakeOutputsAvailable

    -- | No, don't make output available
    --
    -- Use to generate independent transactions.
  | DontMakeOutputsAvailable

-- | State needed for transaction generation
data GenTrState h a = GenTrState {
      -- | State needed to generate inputs
      _gtsInpState :: GenInpState h a

      -- | Next hash
    , _gtsNextHash :: Int
    }

makeLenses ''GenTrState

-- | Initial 'GenTrState'
initTrState :: Utxo h a       -- ^ Initial UTxO
            -> Int            -- ^ First available hash
            -> GenTrState h a
initTrState utxo nextHash = GenTrState {
      _gtsInpState = initInpState utxo
    , _gtsNextHash = nextHash
    }

-- | Transaction generator
type GenTransaction h a = StateT (GenTrState h a) Gen

-- | Try to generate a transaction
--
-- Fails if no inputs were available.
genTransaction :: Hash h a
               => GenTrParams a
               -> RemoveUsedInputs
               -> MakeOutputsAvailable
               -> Set (Input h a)       -- ^ Inputs to avoid
               -> GenTransaction h a (Maybe (Transaction h a))
genTransaction GenTrParams{..} removeUsedInputs makeOutputsAvailable notThese = do
    numInputs <- lift $ chooseNumInputs
    inputs    <- zoom gtsInpState $
                   genDistinctInputs removeUsedInputs notThese numInputs

    if null inputs
      then return Nothing
      else do
        nextHash   <- gtsNextHash <<+= 1
        numOutputs <- lift $ choose (1, gtpMaxNumOutputs)

        let fee          = gtpEstimateFee numInputs numOutputs
            totalInValue = sum (map (outVal . snd) inputs)

        if totalInValue <= fee
          then return Nothing
          else do
            let gos = initOutState (totalInValue - fee)

            outputs <- lift $ (`evalStateT` gos) $
              replicateAtMostM numOutputs $ genOutput gtpOutParams

            let tr = Transaction {
                         trFresh = 0
                       , trIns   = Set.fromList (map fst inputs)
                       , trOuts  = outputs
                       , trFee   = fee
                       , trHash  = nextHash
                       , trExtra = []
                       }

            case makeOutputsAvailable of
              DontMakeOutputsAvailable -> return ()
              MakeOutputsAvailable     -> zoom gtsInpState $
                                            (gisUtxo %= utxoUnion (trUtxo tr))

            return $ Just tr
  where
    -- Bias towards fewer inputs
    chooseNumInputs :: Gen Int
    chooseNumInputs = frequency $ zip [gtpMaxNumInputs, gtpMaxNumInputs-1 ..]
                                      (map pure [1 .. gtpMaxNumInputs])

{-------------------------------------------------------------------------------
  Chains
-------------------------------------------------------------------------------}

data GenChainParams a = GenChainParams {
      -- | Maximum number of transactions per block
      gcpMaxBlockSize   :: Int

      -- | Maximum number of blocks
    , gcpMaxChainLength :: Int

      -- | Transaction parameters
    , gcpTrParams       :: GenTrParams a
    }

-- | Default 'GenChainParams'
defChainParams :: (Int -> Int -> Value)    -- ^ Fee model
               -> [a]                      -- ^ Address we can generate outputs for
               -> GenChainParams a
defChainParams feeModel addresses = GenChainParams {
      gcpMaxBlockSize   = 20
    , gcpMaxChainLength = 10
    , gcpTrParams       = defTrParams feeModel addresses
    }

-- | Generate an arbitrary chain
--
-- The chain will have at least one block, but blocks may be empty.
genChain :: forall h a. Hash h a
         => GenChainParams a -> GenTransaction h a (Chain h a)
genChain GenChainParams{..} = goChain
  where
    goChain :: GenTransaction h a (Chain h a)
    goChain = OldestFirst <$> do
        chainLength <- lift $ choose (1, gcpMaxChainLength)
        replicateM chainLength goBlock

    goBlock :: GenTransaction h a (Block h a)
    goBlock = OldestFirst <$> do
        blockSize <- lift $ choose (0, gcpMaxBlockSize)
        replicateAtMostM blockSize $
          genTransaction gcpTrParams RemoveUsedInputs MakeOutputsAvailable Set.empty

{-------------------------------------------------------------------------------
  Auxiliary
-------------------------------------------------------------------------------}

replicateAtMostM :: Monad m => Int -> m (Maybe a) -> m [a]
replicateAtMostM 0 _ = return []
replicateAtMostM n f = do
    ma <- f
    case ma of
      Just a  -> (a :) <$> replicateAtMostM (n - 1) f
      Nothing -> return []
