--------------------------------------------------------------------------------
-- | This is a small pretty-printing library.
{-# LANGUAGE DeriveFoldable             #-}
{-# LANGUAGE DeriveFunctor              #-}
{-# LANGUAGE DeriveTraversable          #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE RecordWildCards            #-}
module Patat.PrettyPrint.Internal
    ( Control (..)
    , Chunk (..)
    , Chunks
    , chunkToString
    , chunkLines

    , DocE (..)
    , chunkToDocE

    , Indentation (..)

    , Doc (..)
    , docToChunks

    , toString
    , dimensions
    , null

    , hPutDoc
    , putDoc
    , mkDoc
    , string
    ) where


--------------------------------------------------------------------------------
import           Control.Monad.Reader       (asks, local)
import           Control.Monad.RWS          (RWS, runRWS)
import           Control.Monad.State        (get, modify, state)
import           Control.Monad.Writer       (tell)
import           Data.Char.WCWidth.Extended (wcstrwidth)
import qualified Data.List                  as L
import           Data.String                (IsString (..))
import           Prelude                    hiding (null)
import qualified System.Console.ANSI        as Ansi
import qualified System.IO                  as IO


--------------------------------------------------------------------------------
-- | Control actions for the terminal.
data Control
    = ClearScreenControl
    | GoToLineControl Int
    deriving (Eq, Show)


--------------------------------------------------------------------------------
-- | Hyperlinks can be printed in multiple parts when formatting is used.  The
-- ID is used to tell which ones belong to the same link.
type HyperlinkID = Int


--------------------------------------------------------------------------------
-- | A simple chunk of text.  All ANSI codes are "reset" after printing.
data Chunk
    = StringChunk [Ansi.SGR] String
    | HyperlinkChunk HyperlinkID [Ansi.SGR] String String -- Codes, title, URL
    | NewlineChunk
    | ControlChunk Control
    deriving (Eq, Show)


--------------------------------------------------------------------------------
type Chunks = [Chunk]


--------------------------------------------------------------------------------
hPutChunk :: IO.Handle -> Chunk -> IO ()
hPutChunk h NewlineChunk = IO.hPutStrLn h ""
hPutChunk h (StringChunk codes str) = do
    Ansi.hSetSGR h (reverse codes)
    IO.hPutStr h str
    Ansi.hSetSGR h [Ansi.Reset]
hPutChunk h (HyperlinkChunk hid codes str url) = do
    Ansi.hSetSGR h (reverse codes)
    Ansi.hHyperlinkWithId h (show hid) url str
    Ansi.hSetSGR h [Ansi.Reset]
hPutChunk h (ControlChunk ctrl) = case ctrl of
    ClearScreenControl -> Ansi.hClearScreen h
    GoToLineControl l  -> Ansi.hSetCursorPosition h l 0


--------------------------------------------------------------------------------
chunkToString :: Chunk -> String
chunkToString NewlineChunk               = "\n"
chunkToString (StringChunk _ str)        = str
chunkToString (HyperlinkChunk _ _ str _) = str
chunkToString (ControlChunk _)           = ""


--------------------------------------------------------------------------------
-- | If two neighboring chunks have the same set of ANSI codes, we can group
-- them together.
optimizeChunks :: Chunks -> Chunks
optimizeChunks chunks = case chunks of
    (StringChunk c1 s1 : StringChunk c2 s2 : t)
        | c1 == c2 -> optimizeChunks (StringChunk c1 (s1 <> s2) : t)
        | otherwise ->
            StringChunk c1 s1 : optimizeChunks (StringChunk c2 s2 : t)
    (HyperlinkChunk i1 c1 s1 u1 : HyperlinkChunk i2 c2 s2 u2 : t)
        | i1 == i2 && c1 == c2 && u1 == u2 -> optimizeChunks $
            HyperlinkChunk i1 c1 (s1 <> s2) u1 : t
        | otherwise ->
            HyperlinkChunk i1 c1 s1 u1 :
            optimizeChunks (HyperlinkChunk i2 c2 s2 u2 : t)
    x : t -> x : optimizeChunks t
    [] -> []


--------------------------------------------------------------------------------
chunkLines :: Chunks -> [Chunks]
chunkLines chunks = case break (== NewlineChunk) chunks of
    (xs, _newline : ys) -> xs : chunkLines ys
    (xs, [])            -> [xs]


--------------------------------------------------------------------------------
data DocE d
    = String String
    | Softspace
    | Hardspace
    | Softline
    | Hardline
    | WrapAt
        { wrapAtCol :: Maybe Int
        , wrapDoc   :: d
        }
    | Ansi
        { ansiCode :: [Ansi.SGR] -> [Ansi.SGR]  -- ^ Modifies current codes.
        , ansiDoc  :: d
        }
    | Indent
        { indentFirstLine  :: Indentation [Chunk]
        , indentOtherLines :: Indentation [Chunk]
        , indentDoc        :: d
        }
    | Hyperlink String d
    | Control Control
    deriving (Functor)


--------------------------------------------------------------------------------
chunkToDocE :: Chunk -> DocE Doc
chunkToDocE chunk = case chunk of
    NewlineChunk              -> Hardline
    ControlChunk ctrl         -> Control ctrl
    StringChunk c1 s          -> Ansi (\c0 -> c1 ++ c0) $ Doc [String s]
    HyperlinkChunk _ c1 s url -> Ansi (\c0 -> c1 ++ c0) $ Doc
        [Hyperlink url $ Doc [String s]]


--------------------------------------------------------------------------------
newtype Doc = Doc {unDoc :: [DocE Doc]}
    deriving (Monoid, Semigroup)


--------------------------------------------------------------------------------
instance Show Doc where
    show = toString


--------------------------------------------------------------------------------
instance IsString Doc where
    fromString = string


--------------------------------------------------------------------------------
data DocEnv = DocEnv
    { -- | Most recent ones first
      deCodes     :: [Ansi.SGR]
    , -- | First-line indent not included
      deIndent    :: [Indentation [Chunk]]
    , -- | Wrap at columns
      deWrap      :: Maybe Int
    , -- | Hyperlink context
      deHyperlink :: Maybe (HyperlinkID, String)
    }


--------------------------------------------------------------------------------
data DocState = DocState
    { dsLineBuffer  :: LineBuffer
    , dsHyperlinkID :: Int
    }


--------------------------------------------------------------------------------
type DocM = RWS DocEnv Chunks DocState


--------------------------------------------------------------------------------
-- | Note that the lists here are reversed so we have fast append.
-- We also store the current length to avoid having to recompute it.
data LineBuffer = LineBuffer Int [Indentation [Chunk]] [Chunk]


--------------------------------------------------------------------------------
emptyLineBuffer :: LineBuffer
emptyLineBuffer = LineBuffer 0 [] []


--------------------------------------------------------------------------------
data Indentation a = Indentation Int a
    deriving (Foldable, Functor, Traversable)


--------------------------------------------------------------------------------
indentationToChunks :: Indentation [Chunk] -> [Chunk]
indentationToChunks (Indentation 0 c) = c
indentationToChunks (Indentation n c) = StringChunk [] (replicate n ' ') : c


--------------------------------------------------------------------------------
indentationWidth :: Indentation [Chunk] -> Int
indentationWidth (Indentation s c) =
    s + sum (map (wcstrwidth . chunkToString) c)


--------------------------------------------------------------------------------
bufferToChunks :: LineBuffer -> Chunks
bufferToChunks (LineBuffer _ ind chunks) = case chunks of
    [] -> concatMap indentationToChunks $ reverse $
        dropWhile emptyIndentation ind
    _ -> concatMap indentationToChunks (reverse ind) ++ reverse chunks
  where
    emptyIndentation (Indentation _ []) = True
    emptyIndentation _                  = False


--------------------------------------------------------------------------------
docToChunks :: Doc -> Chunks
docToChunks doc0 =
    let env0        = DocEnv [] [] Nothing Nothing
        state0      = DocState emptyLineBuffer 0
        ((), finalState, cs) = runRWS (go $ unDoc doc0) env0 state0 in
    optimizeChunks (cs <> bufferToChunks (dsLineBuffer finalState))
  where
    go :: [DocE Doc] -> DocM ()

    go [] = return ()

    go (String str : docs) = do
        chunk <- makeChunk str
        appendChunk chunk
        go docs

    go (Hyperlink url doc : docs) = do
        hid <- state $ \s ->
            let hid = dsHyperlinkID s in (hid, s {dsHyperlinkID = hid + 1})
        local (\env -> env {deHyperlink = Just (hid, url)}) (go $ unDoc doc)
        go docs

    go (Softspace : docs) = do
        hard <- softConversion Softspace docs
        go (hard : docs)

    go (Hardspace : docs) = do
        chunk <- makeChunk " "
        appendChunk chunk
        go docs

    go (Softline : docs) = do
        hard <- softConversion Softline docs
        go (hard : docs)

    go (Hardline : docs) = do
        buffer <- dsLineBuffer <$> get
        tell $ bufferToChunks buffer <> [NewlineChunk]
        ind <- asks deIndent
        let indw = sum $ map indentationWidth ind
        modify $ \s -> case docs of
            []    -> s {dsLineBuffer = emptyLineBuffer}
            _ : _ -> s {dsLineBuffer = LineBuffer indw ind []}
        go docs

    go (WrapAt {..} : docs) = do
        il <- asks $ sum . map indentationWidth . deIndent
        local (\env -> env {deWrap = fmap (+ il) wrapAtCol}) $ go (unDoc wrapDoc)
        go docs

    go (Ansi {..} : docs) = do
        local (\env -> env {deCodes = ansiCode (deCodes env)}) $
            go (unDoc ansiDoc)
        go docs

    go (Indent {..} : docs) = do
        local (\e -> e {deIndent = indentOtherLines : deIndent e}) $ do
            modify $ \s ->
                let LineBuffer w i c = dsLineBuffer s
                    w' = w + indentationWidth indentFirstLine in
                s {dsLineBuffer = LineBuffer w' (indentFirstLine : i) c}
            go (unDoc indentDoc)
        go docs

    go (Control ctrl : docs) = do
        tell [ControlChunk ctrl]
        go docs

    makeChunk :: String -> DocM Chunk
    makeChunk str = do
        codes <- asks deCodes
        mbHyperlink <- asks deHyperlink
        return $ case mbHyperlink of
            Just (hid, url) -> HyperlinkChunk hid codes str url
            Nothing         -> StringChunk codes str

    appendChunk :: Chunk -> DocM ()
    appendChunk c = modify $ \s ->
        let LineBuffer w i cs = dsLineBuffer s
            w' = w + wcstrwidth (chunkToString c) in
        s {dsLineBuffer = LineBuffer w' i (c : cs)}

    -- Convert 'Softspace' or 'Softline' to 'Hardspace' or 'Hardline'
    softConversion :: DocE Doc -> [DocE Doc] -> DocM (DocE Doc)
    softConversion soft docs = do
        mbWrapCol <- asks deWrap
        case mbWrapCol of
            Nothing     -> return hard
            Just maxCol -> do
                LineBuffer currentCol _ _ <- dsLineBuffer <$> get
                case nextWordLength docs of
                    Nothing                            -> return hard
                    Just l
                        | currentCol + 1 + l <= maxCol -> return Hardspace
                        | otherwise                    -> return Hardline
      where
        hard = case soft of
            Softspace -> Hardspace
            Softline  -> Hardline
            _         -> soft

    nextWordLength :: [DocE Doc] -> Maybe Int
    nextWordLength []                  = Nothing
    nextWordLength (String x : xs)
        | L.null x                     = nextWordLength xs
        | otherwise                    = Just (wcstrwidth x)
    nextWordLength (Hyperlink t _ : _) = Just (wcstrwidth t)
    nextWordLength (Softspace : xs)    = nextWordLength xs
    nextWordLength (Hardspace : xs)    = nextWordLength xs
    nextWordLength (Softline : xs)     = nextWordLength xs
    nextWordLength (Hardline : _)      = Nothing
    nextWordLength (WrapAt {..} : xs)  = nextWordLength (unDoc wrapDoc   ++ xs)
    nextWordLength (Ansi   {..} : xs)  = nextWordLength (unDoc ansiDoc   ++ xs)
    nextWordLength (Indent {..} : xs)  = nextWordLength (unDoc indentDoc ++ xs)
    nextWordLength (Control _ : _)     = Nothing


--------------------------------------------------------------------------------
toString :: Doc -> String
toString = concat . map chunkToString . docToChunks


--------------------------------------------------------------------------------
-- | Returns the rows and columns necessary to render this document
dimensions :: Doc -> (Int, Int)
dimensions doc =
    let ls = lines (toString doc) in
    (length ls, foldr max 0 (map wcstrwidth ls))


--------------------------------------------------------------------------------
null :: Doc -> Bool
null doc = case unDoc doc of [] -> True; _ -> False


--------------------------------------------------------------------------------
hPutDoc :: IO.Handle -> Doc -> IO ()
hPutDoc h = mapM_ (hPutChunk h) . docToChunks


--------------------------------------------------------------------------------
putDoc :: Doc -> IO ()
putDoc = hPutDoc IO.stdout


--------------------------------------------------------------------------------
mkDoc :: DocE Doc -> Doc
mkDoc e = Doc [e]


--------------------------------------------------------------------------------
string :: String -> Doc
string ""  = Doc []
string str = mkDoc $ String str  -- TODO (jaspervdj): Newline conversion?