module Pretty(PrTree(..), Pretty(..), pprint, pprintOneline) where

import Data.Char
import Data.List
import Data.Ord


data PrTree = Leaf String
            | Node String [PrTree]
            -- Bracket "(" ")" "," [1,2,3] => (1, 2, 3)
            | Bracket String String String [PrTree]

class Pretty a where
    pretty :: a -> PrTree

data Budget = Budget { bNode :: !Int, bBracket :: !Int }

decNode :: Budget -> Budget
decNode b = b { bNode = bNode b - 1 }

decBracket :: Budget -> Budget
decBracket b = b { bBracket = bBracket b - 1 }

printingBudget :: Budget
printingBudget = Budget 5 5

pprint :: Pretty a => Int -> a -> String
pprint w val = pprintPr printingBudget w (pretty val)

pprintOneline :: Pretty a => a -> String
pprintOneline = pprint maxBound

pprintPr :: Budget -> Int -> PrTree -> String
pprintPr budget w val = intercalate "\n" $ pprintPrX budget w 0 val

pprintPrOneline :: Budget -> PrTree -> String
pprintPrOneline budget = pprintPr budget maxBound

pprintPrX :: Budget -> Int -> Int -> PrTree -> [String]
pprintPrX _      _   _ (Leaf str) = [str]
pprintPrX budget wid x (Node "" subs) = map (pprintPrOneline (decNode budget)) subs
pprintPrX budget wid x (Node prefix subs) =
    let subrenders = map (pprintPrX (decNode budget) wid (x + indentWid)) subs
    in if any ((> 1) . length) subrenders
           then prefix : map (spaces indentWid ++) (concat subrenders)
           else [prefix ++ " " ++ intercalate " " (concat subrenders)]
-- pprintPrX budget wid x (Node prefix subs) =
--     let prefix' = case prefix of
--                       "" -> ""
--                       _  -> prefix ++ " "
--     in chooseBest (bNode budget) wid
--         [[prefix' ++ intercalate " " (map (pprintPrOneline (decNode budget)) subs)]
--         ,let subrender = concatMap (pprintPrX (decNode budget) wid (x + length prefix')) subs
--          in case subrender of
--                 [] -> [prefix']
--                 (ln:lns) -> (prefix' ++ ln) : [spaces (length prefix') ++ l | l <- lns]
--         ,let subrender = concatMap (pprintPrX (decNode budget) wid (x + indentWid)) subs
--          in prefix : map (spaces indentWid ++) subrender]
pprintPrX budget wid x (Bracket beg end sep subs) = chooseBest (bBracket budget) wid
    [[beg ++ intercalate (sep ++ " ") (map (pprintPrOneline (decBracket budget)) subs) ++ end]
    ,let subrender = concatMap (fmapLast (++ sep) . pprintPrX (decBracket budget) wid (x + length beg)) subs
     in case subrender of
            [] -> [beg ++ end]
            [ln] -> [beg ++ ln ++ end]
            (ln:lns) -> (beg ++ ln) : fmapLast (++ end) [spaces (length beg) ++ l | l <- lns]
    ,let subrender = concatMap (fmapLast (++ sep) . pprintPrX (decBracket budget) wid (x + indentWid)) subs
     in case subrender of
            [] -> [beg ++ end]
            lns -> [trimRight beg] ++ [spaces indentWid ++ l | l <- lns] ++ [trimLeft end]]

chooseBest :: Int -> Int -> [[String]] -> [String]
chooseBest n _ | n <= 0 = head
chooseBest _ wid = argmin $ \lns -> (sum [max 0 (length l - wid) | l <- lns]
                                    ,length lns)

indentWid :: Int
indentWid = 4

argmin :: (Show a, Show b, Ord b) => (a -> b) -> [a] -> a
argmin f l = fst $ minimumBy (comparing snd) [(x, f x) | x <- l]

spaces :: Int -> String
spaces n = replicate n ' '

fmapLast :: (a -> a) -> [a] -> [a]
fmapLast _ [] = []
fmapLast f [x] = [f x]
fmapLast f (x:y:zs) = x : fmapLast f (y:zs)

trimLeft :: String -> String
trimLeft = dropWhile isSpace

trimRight :: String -> String
trimRight = reverse . trimLeft . reverse