Dice Roller

diceroller.hs

module Main (main) where

import Control.Applicative       (optional)
import Control.Monad             (replicateM)
import Control.Monad.Trans.State (State, runState, state)
import Data.List                 (sortOn)
import Data.Maybe                (isJust)
import Data.Ord                  (Down(Down))
import System.Environment        (getArgs)
import System.Random             (RandomGen, getStdGen, randomR, setStdGen)
import Text.Printf               (printf)

import qualified Text.Parsec          as P
import qualified Text.Parsec.Language as PL
import qualified Text.Parsec.Token    as PT

data Perfect = NormalTens | DoubleTens deriving (Eq, Show)

data RollSpec = RollSpec
  { nDice      :: Int
  , difficulty :: Int
  , perfect    :: Perfect
  } deriving Show

data Categorised = Categorised
  { successes :: [Int]
  , fails     :: [Int]
  , botches   :: [Int]
  , extras    :: Int
  } deriving Show

data Result
  = Success Int
  | Failure
  | Botch Int
    deriving (Show, Eq)

parser :: P.Parsec String () RollSpec
parser = do
  let l = PT.makeTokenParser PL.haskellDef
  n <- fromIntegral <$> PT.decimal l P.<?> "n-dice (number)"
  PT.whiteSpace l
  diff <- fromIntegral <$> PT.decimal l P.<?> "difficulty (number)"
  PT.whiteSpace l
  perf <- optional (PT.reserved l "r") P.<?> "r (10s count twice)"
  P.eof
  pure (RollSpec n diff (if isJust perf then DoubleTens else NormalTens))

categorise :: RollSpec -> [Int] -> Categorised
categorise (RollSpec _n diff perf) =
  foldr go (Categorised [] [] [] 0) . sortOn Down
  where
    go i c@(Categorised s f b e)
      | i == 10 && perf == DoubleTens = c{ successes = i:s, extras = e+1 }
      | i >= diff                     = c{ successes = i:s }
      | i == 1                        = c{ botches   = i:b }
      | otherwise                     = c{ fails     = i:f }

displayCategorised :: Categorised -> String
displayCategorised (Categorised s f b _e) =
  printf "%s | %s | %s"
    (if null s then "-" else unwords (map show s))
    (if null f then "-" else unwords (map show f))
    (if null b then "-" else unwords (map show b))

makeRolls :: RandomGen g => RollSpec -> State g (Result, Categorised)
makeRolls spec = do
  rolls <- replicateM (nDice spec) (state (randomR (1,10)))
  let cat@(Categorised s _f b e) = categorise spec rolls
  let result = case (s, b) of
        ([], []) -> Failure
        ([],  _) -> Botch (length b)
        ( _,  _) | length b < length s ->
                     Success ((length s + e - length b) `max` 0)
                 | otherwise ->
                     Failure
  return (result, cat)

main :: IO ()
main = do
  args <- unwords <$> getArgs
  case P.parse parser "" args of
    Left err -> do
      putStrLn "Usage: <N> <DIFF> [r]"
      print err
    Right spec -> do
      printf "Roll %d dice at difficulty %d%s.\n"
        (nDice spec) (difficulty spec)
        (if perfect spec == DoubleTens
           then ", with tens counting as two successes" else "")
      rand <- getStdGen
      let ((total, cat), rand') = runState (makeRolls spec) rand
      setStdGen rand'
      putStrLn (show total ++ ": " ++ displayCategorised cat)