module Layout where

data Token =
    Token String
    | TokenVLCurly String !Int
    | TokenNL !Int
    deriving(Show)

data Context = NoLayout String | Layout String !Int

layout :: [Token] -> [Context] -> [Token]
layout (TokenNL n:Token "in":rs) (Layout "let" n':ls) = rbrace:Token "in":layout rs ls
layout (TokenNL n:rs) (Layout h n':ls)
    | n == n' = semi:layout rs (Layout h n':ls)
    | n > n' = layout rs (Layout h n':ls)
    | n < n' = rbrace:layout (TokenNL n:rs) ls
layout (TokenNL _:rs) ls = layout rs ls
layout (TokenVLCurly h n:rs) (Layout h' n':ls)
    | n > n' = lbrace:layout rs (Layout h n:Layout h' n':ls)
    | otherwise = lbrace : rbrace : layout rs (Layout h' n':ls)
layout (TokenVLCurly h n:rs) ls = lbrace:layout rs (Layout h n:ls)
layout (t@(Token s):rs) ls | s `elem` fsts layoutBrackets = t:layout rs (NoLayout s:ls)
layout (t@(Token s):rs) ls | s `elem` snds layoutBrackets = case ls of
    Layout _ _:ls -> rbrace:layout (t:rs) ls
    NoLayout _:ls -> t:layout rs ls
    [] -> error $ "unexpected " ++ show s
layout (t@(Token "in"):rs) ls = case ls of
    Layout "let" n:ls -> rbrace:t:layout rs ls
    Layout _ _:ls -> rbrace:layout (t:rs) ls
    ls -> t:layout rs ls
layout (t@(Token ","):rs) (Layout {}:ls) | topNoLayout  ls `elem` [Just "(",Just "["] = rbrace:layout (t:rs) ls
layout (t@(Token "where"):rs) ls = case ls of
    Layout l n : rest | l `notElem` ["where","let","of"] 
        -> rbrace : t : layout rs rest -- 'where' closes 'do'
    _otherwise -> t : layout rs ls

layout (t:rs) ls = t:layout rs ls
layout [] (Layout _ n:ls) = rbrace:layout [] ls
layout [] [] = []

topNoLayout :: [Context] -> Maybe String
topNoLayout [] = Nothing
topNoLayout (NoLayout s:_) = Just s
topNoLayout (_:ls) = topNoLayout ls


semi = Token ";"
lbrace = Token "{"
rbrace = Token "}"

fsts = map fst
snds = map snd

layoutBrackets = [
    ("case","of"),
    ("if","then"),
    ("(",")"),
    ("[","]"),
    ("{","}")
    ]