module Homework.Ch01.Hanoi where

import Data.Maybe

-- | Move pegs from the first peg to the last peg.
-- The moves that were made are returned, but an error is returned if a move
-- can't be made.
hanoi :: (MonadFail m) => [String] -> Int -> m [Move]
hanoi labels numDiscs = do
  -- CONSTRUCT a set of pegs given the provided arguments
  pegsStart <- initPegs labels numDiscs
  -- Make a move
  (moveMade, _) <- runPegs move pegsStart
  -- Cheat the return for now, assume that movesMade is present for TDD
  return [fromJust moveMade]

{------------------------------------------------------------------------------}
{- MAKE A MOVE ----------------------------------------------------------------}
{------------------------------------------------------------------------------}

move :: (MonadFail m) => PegStep m (Maybe Move)
move = do
  topDiscA <- getTopDisc <$> getPeg 0
  topDiscC <- getTopDisc <$> getPeg 2
  if topDiscA >= topDiscC
    then return Nothing
    else do Just <$> makeMove 0 2

{------------------------------------------------------------------------------}
{- PEGS -----------------------------------------------------------------------}
{------------------------------------------------------------------------------}

-- A set of pegs ordered from start to finish.
data Pegs = Pegs
  { pegsList :: [Peg],
    pegsMoves :: [Move]
  }
  deriving (Eq, Show)

-- CONSTRUCT a set of pegs with their labels and number of discs to fill the first peg with
initPegs :: (MonadFail m) => [String] -> Int -> m Pegs
initPegs (firstLabel : restLabels) numDiscs
  | numDiscs > 0 =
    return
      Pegs
        { pegsList = fillPeg firstLabel numDiscs : (emptyPeg <$> restLabels),
          pegsMoves = []
        }
  | otherwise = fail "Can't create pegs with no discs"
initPegs [] _ = fail "Can't create pegs without peg labels"

{------------------------------------------------------------------------------}
{- PEG STEP -------------------------------------------------------------------}
{------------------------------------------------------------------------------}

-- A peg step holds a function that can be run against some pegs
newtype PegStep m a = PegStep {runPegs :: Pegs -> m (a, Pegs)}

-- CONSTRUCT a new PegStep given a function
withPegs :: (Pegs -> m (a, Pegs)) -> PegStep m a
withPegs = PegStep

-- CONSTRUCT a new PegStep which produces the current pegs
getPegs :: (Monad m) => PegStep m Pegs
getPegs = withPegs $ \pegs -> return (pegs, pegs)

-- CONSTRUCT a new PegStep which takes a pegs and replaces the current pegs
putPegs :: (Monad m) => Pegs -> PegStep m ()
putPegs pegs = withPegs $ return . const ((), pegs)

-- GET or ASK for a value from pegs using a function
askPegs :: (Monad m) => (Pegs -> a) -> PegStep m a
askPegs f = f <$> getPegs

-- MODIFY pegs using a function
modifyPegs :: (Monad m) => (Pegs -> PegStep m Pegs) -> PegStep m ()
modifyPegs f = putPegs =<< f =<< getPegs

-- EMBEDDED PEG ACCESSORS ------------------------------------------------------

getPeg :: (MonadFail m) => Int -> PegStep m Peg
getPeg n = do
  list <- askPegs pegsList
  go n (length list) list
  where
    go n' limit (x : xs)
      | n' == 0 = return x
      | otherwise = go (n' - 1) limit xs
    go _ limit [] = fail $ "Peg out of bounds: queried for index " ++ show n ++ " out of " ++ show limit ++ " pegs"

putPeg :: (MonadFail m) => Peg -> PegStep m ()
putPeg peg@Peg {pegLabel = label} = do
  list <- askPegs pegsList
  modifyPegs $ \pegs' -> do
    newList <- go [] list
    return $ pegs' {pegsList = newList}
  where
    go previous (current : next)
      | pegLabel current == label = return $ reverse previous ++ peg : next
      | otherwise = go (current : previous) next
    go _ [] = fail $ "Could not replace peg with label " ++ show label ++ ": peg does not exist"

-- EMBEDDED PEG METHODS --------------------------------------------------------

makeMove :: (MonadFail m) => Int -> Int -> PegStep m Move
makeMove from to = do
  pushPeg to =<< popPeg from
  fromLabel <- pegLabel <$> getPeg from
  toLabel <- pegLabel <$> getPeg to
  let move' = Move fromLabel toLabel
  pushMove move'
  return move'

popPeg :: (MonadFail m) => Int -> PegStep m Disc
popPeg n = do
  peg@Peg {pegLabel = label, pegDiscs = discs} <- getPeg n
  case discs of
    topDisc : rest -> do
      putPeg peg {pegDiscs = rest}
      return topDisc
    _ -> fail $ "Could not pop empty peg " ++ show label

pushPeg :: (MonadFail m) => Int -> Disc -> PegStep m ()
pushPeg n disc = do
  peg <- getPeg n
  case pegDiscs peg of
    discs@(topDisc : _)
      | disc < topDisc -> putPeg peg {pegDiscs = disc : discs}
      | otherwise -> fail $ "Tried to put bigger disc " ++ show disc ++ " on top of peg " ++ show (pegLabel peg) ++ "'s top disc " ++ show topDisc
    _ -> putPeg peg {pegDiscs = [disc]}

pushMove :: (Monad m) => Move -> PegStep m ()
pushMove move' = do
  pegs <- getPegs
  let moves = pegsMoves pegs
  putPegs pegs {pegsMoves = move' : moves}

{------------------------------------------------------------------------------}
{- PEG STEP INSTANCES ---------------------------------------------------------}
{------------------------------------------------------------------------------}

instance (Monad m) => Functor (PegStep m) where
  fn `fmap` step = withPegs $ \startPegs -> do
    (result, resultPegs) <- runPegs step startPegs
    return (fn result, resultPegs)

instance (Monad m) => Applicative (PegStep m) where
  pure x = withPegs $ \pegs -> return (x, pegs)
  fn <*> x = withPegs $ \startPegs -> do
    (fn', middleState) <- runPegs fn startPegs
    (x', resultState) <- runPegs x middleState
    return (fn' x', resultState)

instance (Monad m) => Monad (PegStep m) where
  firstStep >>= secondStepFactory =
    PegStep $ \startPegs -> do
      (firstResult, middlePegs) <- runPegs firstStep startPegs
      runPegs (secondStepFactory firstResult) middlePegs

instance (MonadFail m) => MonadFail (PegStep m) where
  fail message = withPegs $ \_ -> fail message

{------------------------------------------------------------------------------}
{- PEG ------------------------------------------------------------------------}
{------------------------------------------------------------------------------}

-- A peg is labeled and contains a stack of discs
data Peg = Peg {pegLabel :: String, pegDiscs :: [Disc]} deriving (Eq, Show)

-- CONSTRUCT a new peg with a label and number of disks to fill it with
fillPeg :: String -> Int -> Peg
fillPeg label numDiscs =
  Peg
    { pegLabel = label,
      pegDiscs = stackDiscs numDiscs
    }

-- CONSTRUCT an empty peg with a label
emptyPeg :: String -> Peg
emptyPeg label = Peg {pegLabel = label, pegDiscs = []}

-- GET the top disc from the peg, if it exists
getTopDisc :: Peg -> Maybe Disc
getTopDisc = headOption . pegDiscs
  where
    headOption (x : _) = Just x
    headOption [] = Nothing

{------------------------------------------------------------------------------}
{- DISC -----------------------------------------------------------------------}
{------------------------------------------------------------------------------}

-- A Disc has a size.
data Disc = Disc {discSize :: Int} deriving (Eq, Ord, Show)

-- CONSTRUCT a stack of discs
stackDiscs :: Int -> [Disc]
stackDiscs numDiscs = Disc <$> [1 .. numDiscs]

{------------------------------------------------------------------------------}
{- MOVE -----------------------------------------------------------------------}
{------------------------------------------------------------------------------}

-- A move has the peg that the disc was moved from and the peg it was moved to
data Move = Move {moveFrom :: String, moveTo :: String} deriving (Eq, Show)