{- | Annotating tree structures with desired properties for rooms.  -}
module Dodge.Annotation 
    ( module Dodge.Annotation.Data
    , module Dodge.Annotation
    ) where
import Dodge.RandomHelp
import Dodge.Layout.Tree.Polymorphic
import Dodge.Layout.Tree.Either
import Dodge.Room
import Dodge.Annotation.Data

import Data.Tree
import Control.Monad.State
import System.Random
import Control.Lens

addLock :: RandomGen g => Int -> Tree [Annotation g] -> State g (Tree [Annotation g])
addLock i t = do
    (beforeLock, afterLock) <- splitTrunk t
    newBefore <- applyToRandomNode (Key i :) beforeLock
    return $ addToTrunk newBefore [Node [Lock i] afterLock]
{- | Add one corridor between each parent-child link of a tree of annotations.  -}
padWithCorridors :: Tree [Annotation g] -> Tree [Annotation g]
padWithCorridors = padWithAno [Corridor]

padWithAno :: [Annotation g] -> Tree [Annotation g] -> Tree [Annotation g]
padWithAno ano (Node x xs) = Node ano [Node x (map (padWithAno ano) xs)]

padSucWithCorridors :: Tree [Annotation g] -> Tree [Annotation g]
padSucWithCorridors (Node x xs) = Node x (map padWithCorridors xs)

{- Add one to three corridors between each parent-child link of a tree of annotations.  -}
randomPadCorridors :: RandomGen g => Tree [Annotation g] -> State g (Tree [Annotation g])
randomPadCorridors  (Node x xs) = do
    n <- state $ randomR (1, 3)
    xs' <- mapM randomPadCorridors xs
    return $ treeFromTrunk (replicate n [Corridor]) (Node x xs')
{- | Add a corridor to a random out-link of a room.  -}
roomThenCorridor :: RandomGen g => Room -> State g (Tree (Either Room Room))
roomThenCorridor theRoom = fmap (\r -> Node (Left theRoom) [(pure . Right) r]) 
                                (randomiseOutLinks corridor)

{- | Create a random room tree structure from a list of annotations.  -}
anoToRoomTree :: RandomGen g => [Annotation g] -> State g (Tree (Either Room Room))
anoToRoomTree anos = case anos of
    [AnoApplyInt i f] -> f i
    [SetLabel i randrm] -> do
        rm <- randrm
        return . connectRoom $ rm & rmLabel ?~ i
    [UseLabel i randrm] -> do
        rm <- randrm
        return $ connectRoom $ rm & rmTakeFrom ?~ i
    [ChainAnos ass] -> do
        rms <- mapM anoToRoomTree ass
        return $ linkEitherTrees rms
    [OrAno as] -> do
        a <- takeOne as
        anoToRoomTree a
    [Corridor] -> pure . Right <$> randomiseOutLinks corridor
    [AirlockAno] -> airlock >>=  roomThenCorridor
    [FirstWeapon] -> do
        branchWP <- branchRectWith weaponRoom
        blockedC <- longBlockedCorridor
        join $ takeOne $ return (appendEitherTree branchWP [blockedC]) : replicate 5 weaponRoom
    [EndRoom] -> fmap (pure . Right) (telRoomLev 1)
    (SpecificRoom rt:_) -> rt
    (BossAno cr : _) -> do
        br <- bossRoom cr
        branchRectWith . pure . fmap Left $ treeFromPost [corridor,corridor] br
    (TreasureAno crs loot : _) -> branchRectWith . fmap (pure . Left) $ lootRoom crs loot
    _ -> do
        w <- state $ randomR (100,400)
        h <- state $ randomR (200,400)
        fmap (pure . Right) . randomiseOutLinks $ roomRectAutoLinks w h