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 :: (Monad m) => String -> String -> String -> Int -> m [Move]
hanoi pegLabelA pegLabelB pegLabelC numDiscs = do
  let -- CONSTRUCT a set of pegs given the provided arguments
      pegsStart = initPegs pegLabelA pegLabelB pegLabelC 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 :: (Monad m) => PegStep m (Maybe Move)
move = do
  topDiscA <- getTopDisc <$> getPegA
  topDiscC <- getTopDisc <$> getPegC
  if topDiscA >= topDiscC
    then return Nothing
    else do
      popPegA
      pushPegC $ fromJust topDiscA
      Just <$> makeMove "a" "c"

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

-- A set of pegs ordered from start to finish.
data Pegs = Pegs
  { pegsPegA :: Peg,
    pegsPegB :: Peg,
    pegsPegC :: 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 :: String -> String -> String -> Int -> Pegs
initPegs pegLabelA pegLabelB pegLabelC numDiscs =
  Pegs
    { pegsPegA = fillPeg pegLabelA numDiscs,
      pegsPegB = emptyPeg pegLabelB,
      pegsPegC = emptyPeg pegLabelC,
      pegsMoves = []
    }

{------------------------------------------------------------------------------}
{- 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 -> PegStep m 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 PEGS GETTERS -------------------------------------------------------

getPegA :: (Monad m) => PegStep m Peg
getPegA = askPegs (return . pegsPegA)

putPegA :: (Monad m) => Peg -> PegStep m ()
putPegA peg = do
  pegs <- getPegs
  putPegs $ pegs {pegsPegA = peg}

getPegC :: (Monad m) => PegStep m Peg
getPegC = askPegs (return . pegsPegC)

putPegC :: (Monad m) => Peg -> PegStep m ()
putPegC peg = do
  pegs <- getPegs
  putPegs $ pegs {pegsPegC = peg}

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

makeMove :: (Monad m) => String -> String -> PegStep m Move
makeMove from to = do
  let move' = Move from to
  pegs <- getPegs
  putPegs pegs {pegsMoves = move' : pegsMoves pegs}
  return move'

popPegA :: (Monad m) => PegStep m ()
popPegA = do
  peg <- getPegA
  putPegA $ peg {pegDiscs = init $ pegDiscs peg}

pushPegC :: (Monad m) => Disc -> PegStep m ()
pushPegC disc = do
  peg <- getPegC
  putPegC $ peg {pegDiscs = pegDiscs peg <> [disc]}

{------------------------------------------------------------------------------}
{- 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 = lastOption . pegDiscs
  where
    lastOption xs = case xs of
      [x] -> Just x
      _ : xs' -> lastOption xs'
      [] -> 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 <$> reverse [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)