{-# LANGUAGE StrictData #-}

module Little.Earley.Internal.Render where

import Little.Earley.Internal.Core
import Little.Earley.Internal.Tree
import Little.Earley.Internal.Pretty

-- * Render parse trees

data Box a = Box
  { height :: Int
  , width :: Int
  , contents :: a
  } deriving (Eq, Show)

data PicTree
  = PTBranch String [(Int, PicTree)]  -- Implicit space between columns
  | PTCenterLine Int PicTree
  | PTEmpty
  deriving (Eq, Show)

-- | Draw a tree in the terminal.
drawTree :: (n -> String) -> (c -> String) -> Tree n t c -> [String]
drawTree showN showC = drawPicTree . formatTree showN showC

ptLeaf :: String -> PicTree
ptLeaf s = PTBranch s []

drawPicTree :: Box PicTree -> [String]
drawPicTree u = draw (height u) [(width u, contents u)] where
  draw 0 _ = error "Should not happen"
  draw h xs
    | h == 1 = replicate (width u) '-' : l : []
    | otherwise = l : draw (h-1) ys
    where (l, ys) = drawLine xs

drawLine :: [(Int, PicTree)] -> (String, [(Int, PicTree)])
drawLine = foldMap drawElem

drawElem :: (Int, PicTree) -> (String, [(Int, PicTree)])
drawElem (w, PTEmpty) = (replicate w ' ', [(w, PTEmpty)])
drawElem (w, PTCenterLine h p) = (center w "|", [(w, if h == 1 then p else PTCenterLine (h-1) p)])
drawElem (w, PTBranch s []) = (center w s, [(w, PTEmpty)])
drawElem (w, PTBranch s us) = (hdr, ps) where
  hdr = pasteCenter s (drawConnector (fmap fst ps))
  ps = zipWith (\w' (w0, p) -> (w' + w0, p)) (replicate qw (dw + 1) ++ repeat dw) us
  totalWidthChildren = sum (fmap fst us)
  extraW = w - totalWidthChildren
  (dw, qw) = extraW `divMod` length us

drawConnector :: [Int] -> String
drawConnector [] = error "Should not be empty"
drawConnector [w] = center w "+"
drawConnector (w : w' : ws) = center_ ' ' '-' w "+" ++ drawConnector' w' ws where
  drawConnector' w0 [] = center_ '-' ' ' w0 "+"
  drawConnector' w0 (w1 : ws1) = center_ '-' '-' w0 "+" ++ drawConnector' w1 ws1

pasteCenter :: [a] -> [a] -> [a]
pasteCenter mid full = xl ++ mid ++ xr where
  nmid = length mid
  nfull = length full
  wl = (nfull - nmid) `div` 2
  (xl, full') = splitAt wl full
  xr = drop nmid full'

center :: Int -> String -> String
center = center_ ' ' ' '

center_ :: a -> a -> Int -> [a] -> [a]
center_ cl cr w s = replicate wl cl ++ s ++ replicate wr cr where
  len = length s
  wl = (w - len) `div` 2
  wr = w - len - wl

formatTree :: (n -> String) -> (c -> String) -> Tree n t c -> Box PicTree
formatTree _drawN drawC (Leaf _ _t c) = Box
  { height = 1
  , width = length c'
  , contents = ptLeaf c'
  } where c' = drawC c
formatTree drawN _drawC (Brch (RuleId n j) _ _ []) = Box
  { height = 3
  , width = length nj'
  , contents = ptLeaf nj'
  } where nj' = drawN n ++ " #" ++ show j
formatTree drawN drawC (Brch (RuleId n j) _ _ us) = Box
  { height = height'
  , width = width'
  , contents = PTBranch nj'
      [(width u', PTCenterLine (maxHeight - height u' + 1) (contents u')) | u' <- us']
  } where
      us' = fmap (formatTree drawN drawC) us
      totalWidthChildren = sum (fmap width us') + length us' - 1
      maxHeight = maximum (fmap height us')
      nj' = drawN n ++ " #" ++ show j
      height' = maxHeight + 2
      width' = max totalWidthChildren (length nj')

-- | 'drawTree' using 'prettyPrint' to show symbols.
prettyTree :: (PrettyPrint n, PrettyPrint c) => Tree n t c -> [String]
prettyTree = drawTree prettyPrint prettyPrint