{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE UndecidableInstances #-}
module Symantic.Document.Plain where

import Data.Bool
import Data.Char (Char)
import Data.Eq (Eq(..))
import Data.Foldable (foldr)
import Data.Function (($), (.), id)
import Data.Functor ((<$>))
import Data.Maybe (Maybe(..))
import Data.Monoid (Monoid(..))
import Data.Ord (Ord(..), Ordering(..))
import Data.Semigroup (Semigroup(..))
import Data.String (String, IsString(..))
import Data.Text (Text)
import GHC.Natural (minusNatural,quotRemNatural)
import Numeric.Natural (Natural)
import Prelude (fromIntegral, Num(..))
import System.Console.ANSI
import Text.Show (Show(..), showString, showParen)
import qualified Data.List as List
import qualified Data.Text.Lazy as TL

import Symantic.Document.API

-- * Type 'Plain'
-- | Church encoded for performance concerns.
-- Kind like 'ParsecT' in @megaparsec@ but a little bit different
-- due to the use of 'PlainFit' for implementing 'breakingSpace' correctly
-- when in the left hand side of ('<>').
-- Prepending is done using continuation, like in a difference list.
newtype Plain d = Plain
 { unPlain ::
     {-curr-}PlainInh ->
     {-curr-}PlainState d ->
     {-ok-}( ({-prepend-}(d->d), {-new-}PlainState d) -> PlainFit d) ->
     PlainFit d
     -- NOTE: equivalent to:
     -- ReaderT PlainInh (StateT (PlainState d) (Cont (PlainFit d))) (d->d)
 }

runPlain :: Monoid d => Plain d -> d
runPlain x =
        unPlain x
         defPlainInh
         defPlainState
         {-k-}(\(px,_sx) fits _overflow ->
                -- NOTE: if px fits, then appending mempty fits
                fits (px mempty) )
         {-fits-}id
         {-overflow-}id

-- ** Type 'PlainState'
data PlainState d = PlainState
 { plainState_buffer          :: ![PlainChunk d]
 , plainState_bufferStart     :: !Column
 , plainState_bufferWidth     :: !Width
 , plainState_removableIndent :: !Indent
   -- ^ The amount of 'Indent' added by 'breakspace'
   -- that can be removed by breaking the 'space' into a 'newline'.
 } deriving (Show)

defPlainState :: PlainState d
defPlainState = PlainState
 { plainState_buffer          = mempty
 , plainState_bufferStart     = 0
 , plainState_bufferWidth     = 0
 , plainState_removableIndent = 0
 }

-- ** Type 'PlainInh'
data PlainInh = PlainInh
 { plainInh_width   :: !(Maybe Column)
 , plainInh_justify :: !Bool
 , plainInh_indent  :: !Width
 } deriving (Show)

defPlainInh :: PlainInh
defPlainInh = PlainInh
 { plainInh_width   = Nothing
 , plainInh_justify = False
 , plainInh_indent  = 0
 }

-- ** Type 'PlainFit'
-- | Double continuation to qualify the returned document
-- as fitting or overflowing the given 'plainInh_width'.
-- It's like @('Bool',d)@ in a normal style
-- (a non continuation-passing-style).
type PlainFit d = {-fits-}(d -> d) ->
                  {-overflow-}(d -> d) ->
                  d

-- ** Type 'PlainChunk'
data PlainChunk d
 =   PlainChunk_Ignored d
     -- ^ Ignored by the justification but kept in place.
     -- Used for instance to put ANSI sequences.
 |   PlainChunk_Word (Word d)
 |   PlainChunk_Spaces Width
     -- ^ 'spaces' preserved to be interleaved
     -- correctly with 'PlainChunk_Ignored'.
instance Show d => Show (PlainChunk d) where
        showsPrec p x =
                showParen (p>10) $
                case x of
                 PlainChunk_Ignored d ->
                        showString "Z " .
                        showsPrec 11 d
                 PlainChunk_Word (Word d) ->
                        showString "W " .
                        showsPrec 11 d
                 PlainChunk_Spaces s ->
                        showString "S " .
                        showsPrec 11 s
instance Lengthable d => Lengthable (PlainChunk d) where
        length = \case
         PlainChunk_Ignored{} -> 0
         PlainChunk_Word d -> length d
         PlainChunk_Spaces s -> s
        nullLength = \case
         PlainChunk_Ignored{} -> True
         PlainChunk_Word d -> nullLength d
         PlainChunk_Spaces s -> s == 0
instance DocFrom [SGR] d => DocFrom [SGR] (PlainChunk d) where
        docFrom sgr = PlainChunk_Ignored (docFrom sgr)

runPlainChunk :: Spaceable d => PlainChunk d -> d
runPlainChunk = \case
 PlainChunk_Ignored d -> d
 PlainChunk_Word (Word d) -> d
 PlainChunk_Spaces s -> spaces s

instance Semigroup d => Semigroup (Plain d) where
        Plain x <> Plain y = Plain $ \inh st k ->
                x inh st $ \(px,sx) ->
                        y inh sx $ \(py,sy) ->
                                k (px.py,sy)
instance Monoid d => Monoid (Plain d) where
        mempty = Plain $ \_inh st k -> k (id,st)
        mappend = (<>)
instance (Spaceable d) => Spaceable (Plain d) where
        newline = Plain $ \inh st k ->
                k(\next ->
                        (if plainInh_justify inh then joinLine inh st else mempty) <>
                        newline<>spaces (plainInh_indent inh)<>next
                 , st
                         { plainState_bufferStart      = plainInh_indent inh
                         , plainState_bufferWidth      = 0
                         , plainState_buffer           = mempty
                         }
                 )
        space = spaces 1
        spaces n = Plain $ \inh st@PlainState{..} k fits overflow ->
                let newWidth = plainState_bufferStart + plainState_bufferWidth + n in
                if plainInh_justify inh
                then
                        let newState =
                                case plainState_buffer of
                                 PlainChunk_Spaces s:bs -> st
                                         { plainState_buffer      = PlainChunk_Spaces (s+n):bs
                                         }
                                 _ -> st
                                         { plainState_buffer      = PlainChunk_Spaces n:plainState_buffer
                                         , plainState_bufferWidth = plainState_bufferWidth + 1
                                         }
                        in
                        case plainInh_width inh of
                         Just width | width < newWidth ->
                                overflow $ k (id{-(d<>)-}, newState) fits overflow
                         _ -> k (id{-(d<>)-}, newState) fits overflow
                else
                        let newState = st
                                 { plainState_bufferWidth = plainState_bufferWidth + n
                                 } in
                        case plainInh_width inh of
                         Just width | width < newWidth ->
                                overflow $ k ((spaces n <>), newState) fits fits
                         _ -> k ((spaces n <>), newState) fits overflow
instance (DocFrom (Word s) d, Semigroup d, Lengthable s) => DocFrom (Word s) (Plain d) where
        docFrom s = Plain $ \inh st@PlainState{..} k fits overflow ->
                let wordLen = length s in
                if wordLen <= 0
                then k (id,st) fits overflow
                else
                        let newBufferWidth = plainState_bufferWidth + wordLen in
                        let newWidth = plainState_bufferStart + newBufferWidth in
                        if plainInh_justify inh
                        then
                                let newState = st
                                         { plainState_buffer =
                                                PlainChunk_Word (Word (docFrom s)) :
                                                plainState_buffer
                                         , plainState_bufferWidth = newBufferWidth
                                         } in
                                case plainInh_width inh of
                                 Just width | width < newWidth ->
                                        overflow $ k (id, newState) fits overflow
                                 _ -> k (id, newState) fits overflow
                        else
                                let newState = st
                                         { plainState_bufferWidth = newBufferWidth
                                         } in
                                case plainInh_width inh of
                                 Just width | width < newWidth ->
                                        overflow $ k ((docFrom s <>), newState) fits fits
                                 _ -> k ((docFrom s <>), newState) fits overflow
instance (DocFrom (Word s) d, Lengthable s, Spaceable d, Splitable s) =>
         DocFrom (Line s) (Plain d) where
        docFrom =
                mconcat .
                List.intersperse breakspace .
                (docFrom <$>) .
                words .
                unLine
instance Spaceable d => Indentable (Plain d) where
        align p = Plain $ \inh st ->
                let currInd = plainState_bufferStart st + plainState_bufferWidth st in
                unPlain p inh{plainInh_indent=currInd} st
        incrIndent i p = Plain $ \inh ->
                unPlain p inh{plainInh_indent = plainInh_indent inh + i}
        setIndent i p = Plain $ \inh ->
                unPlain p inh{plainInh_indent=i}
        fill m p = Plain $ \inh0 st0 ->
                let col0 = plainState_bufferStart st0 + plainState_bufferWidth st0 in
                let p1 = Plain $ \inh1 st1 ->
                        let col1 = plainState_bufferStart st1 + plainState_bufferWidth st1 in
                        let w | col0 <= col1 = col1`minusNatural`col0
                              | otherwise    = col0`minusNatural`col1 in
                        unPlain
                         (if w<=m
                                then spaces (m`minusNatural`w)
                                else mempty)
                         inh1 st1
                in
                unPlain (p <> p1) inh0 st0
        breakfill m p = Plain $ \inh0 st0 ->
                let col0 = plainState_bufferStart st0 + plainState_bufferWidth st0 in
                let p1 = Plain $ \inh1 st1 ->
                        let col1 = plainState_bufferStart st1 + plainState_bufferWidth st1 in
                        let w | col0 <= col1 = col1`minusNatural`col0
                              | otherwise    = col0`minusNatural`col1 in
                        unPlain
                         (case w`compare`m of
                                 LT -> spaces (m`minusNatural`w)
                                 EQ -> mempty
                                 GT -> setIndent (col0 + m) newline)
                         inh1 st1
                in
                unPlain (p <> p1) inh0 st0
instance Spaceable d => Justifiable (Plain d) where
        justify p = (<> flushLastLine) $ Plain $ \inh ->
                unPlain p inh{plainInh_justify=True}
                where
                flushLastLine :: Plain d
                flushLastLine = Plain $ \_inh st@PlainState{..} ok ->
                        ok
                         ( (joinPlainLine (collapseSpaces <$> List.reverse plainState_buffer) <>)
                         , st
                                 { plainState_bufferStart      = plainState_bufferStart + plainState_bufferWidth
                                 , plainState_bufferWidth      = 0
                                 , plainState_buffer           = mempty
                                 }
                         )
                collapseSpaces = \case
                 PlainChunk_Spaces s -> PlainChunk_Spaces (if s > 0 then 1 else 0)
                 x -> x
instance (Spaceable d) => Wrappable (Plain d) where
        setWidth w p = Plain $ \inh ->
                unPlain p inh{plainInh_width=w}
        breakpoint = Plain $ \inh st k fits overflow ->
                let newlineInd = plainInh_indent inh in
                k
                 ( id
                 , st
                         { plainState_removableIndent  = newlineInd
                         }
                 )
                 fits
                 {-overflow-}(\_r ->
                        unPlain newline inh st k
                         fits
                         {-overflow-}(
                                if plainState_removableIndent st < newlineInd
                                then overflow
                                else fits
                         )
                 )
        breakspace = Plain $ \inh st k fits overflow ->
                let newlineInd = plainInh_indent inh in
                k
                 ( if plainInh_justify inh then id else (space <>)
                 , st
                         { plainState_buffer =
                                case plainState_buffer st of
                                 PlainChunk_Spaces s:bs -> PlainChunk_Spaces (s+1):bs
                                 bs -> PlainChunk_Spaces 1:bs
                         , plainState_bufferWidth      = plainState_bufferWidth st + 1
                         , plainState_removableIndent  = newlineInd
                         }
                 )
                 fits
                 {-overflow-}(\_r ->
                        unPlain newline inh st k
                         fits
                         {-overflow-}(
                                if plainState_removableIndent st < newlineInd
                                then overflow
                                else fits
                         )
                 )
        breakalt x y = Plain $ \inh st k fits overflow ->
                unPlain x inh st k fits
                 {-overflow-}(\_r ->
                        unPlain y inh st k fits overflow
                 )
-- String
instance (DocFrom (Word String) d, Spaceable d) =>
         DocFrom String (Plain d) where
        docFrom =
                mconcat .
                List.intersperse newline .
                (docFrom <$>) .
                lines
instance (DocFrom (Word String) d, Spaceable d) =>
         IsString (Plain d) where
        fromString = docFrom
-- Text
instance (DocFrom (Word Text) d, Spaceable d) =>
         DocFrom Text (Plain d) where
        docFrom =
                mconcat .
                List.intersperse newline .
                (docFrom <$>) .
                lines
instance (DocFrom (Word TL.Text) d, Spaceable d) =>
         DocFrom TL.Text (Plain d) where
        docFrom =
                mconcat .
                List.intersperse newline .
                (docFrom <$>) .
                lines
-- Char
instance (DocFrom (Word Char) d, Spaceable d) =>
         DocFrom Char (Plain d) where
        docFrom ' '  = breakspace
        docFrom '\n' = newline
        docFrom c    = docFrom (Word c)


instance (DocFrom [SGR] d, Semigroup d) => DocFrom [SGR] (Plain d) where
        docFrom sgr = Plain $ \inh st k ->
                if plainInh_justify inh
                then k (id, st {plainState_buffer = PlainChunk_Ignored (docFrom sgr) : plainState_buffer st})
                else k ((docFrom sgr <>), st)

joinLine ::
 Spaceable d =>
 PlainInh -> PlainState d -> d
joinLine PlainInh{..} PlainState{..} =
        case plainInh_width of
         Nothing -> joinPlainLine $ List.reverse plainState_buffer
         Just width ->
                if width < plainState_bufferStart
                || width < plainInh_indent
                then joinPlainLine $ List.reverse plainState_buffer
                else
                        let wordCount =
                                foldr (\c acc -> acc + case c of
                                 PlainChunk_Word{} -> 1
                                 _ -> 0) 0 plainState_buffer in
                        let bufferWidth =
                                -- NOTE: compress all separated spaces into a single one.
                                if wordCount == 0
                                then 0
                                else
                                        let spaceCount = foldr
                                                 (\c acc ->
                                                        acc + case c of
                                                         PlainChunk_Ignored{} -> 0
                                                         PlainChunk_Word{} -> 0
                                                         PlainChunk_Spaces s -> s)
                                                 0 plainState_buffer in
                                        (plainState_bufferWidth`minusNatural`spaceCount) +
                                        (wordCount`minusNatural`1) in
                        let adjustedWidth =
                                max bufferWidth $
                                        min
                                        (width`minusNatural`plainState_bufferStart)
                                        (width`minusNatural`plainInh_indent) in
                        unLine $ padPlainLineInits adjustedWidth
                         (bufferWidth,wordCount,List.reverse plainState_buffer)

-- | @('justifyPadding' a b)@ returns the padding lengths
-- to reach @(a)@ in @(b)@ pads,
-- using the formula: @(a '==' m'*'(q '+' q'+'1) '+' ('r'-'m)'*'(q'+'1) '+' (b'-'r'-'m)'*'q)@
-- where @(q+1)@ and @(q)@ are the two padding lengths used and @(m = min (b-r) r)@.
--
-- A simple implementation of 'justifyPadding' could be:
-- @
-- 'justifyPadding' a b =
--   'join' ('List.replicate' m [q,q'+'1])
--   <> ('List.replicate' (r'-'m) (q'+'1)
--   <> ('List.replicate' ((b'-'r)'-'m) q
--   where
--   (q,r) = a`divMod`b
--   m = 'min' (b-r) r
-- @
justifyPadding :: Natural -> Natural -> [Natural]
justifyPadding a b = go r (b-r) -- NOTE: r >= 0 && b-r >= 0 due to 'divMod'
        where
        (q,r) = a`quotRemNatural`b

        go 0  bmr = List.replicate (fromIntegral bmr) q    -- when min (b-r) r == b-r
        go rr 0   = List.replicate (fromIntegral rr) (q+1) -- when min (b-r) r == r
        go rr bmr = q:(q+1) : go (rr`minusNatural`1) (bmr`minusNatural`1)

padPlainLineInits ::
 Spaceable d =>
 Width -> (Natural, Natural, [PlainChunk d]) -> Line d
padPlainLineInits width (lineLen,wordCount,line) = Line $
        if width <= lineLen
                -- The gathered line reached or overreached the width,
                -- hence no padding id needed.
        || wordCount <= 1
                -- The case width <= lineLen && wordCount == 1
                -- can happen if first word's length is < width
                -- but second word's len is >= width.
        then joinPlainLine line
        else
                -- Share the missing spaces as evenly as possible
                -- between the words of the line.
                padPlainLine line $ justifyPadding (width-lineLen) (wordCount-1)

joinPlainLine :: Monoid d => Spaceable d => [PlainChunk d] -> d
joinPlainLine = mconcat . (runPlainChunk <$>)

padPlainLine :: Spaceable d => [PlainChunk d] -> [Width] -> d
padPlainLine = go
        where
        go (w:ws) lls@(l:ls) =
                case w of
                 PlainChunk_Spaces _s -> spaces (fromIntegral (l+1)) <> go ws ls
                 _ -> runPlainChunk w <> go ws lls
        go (w:ws) [] = runPlainChunk w <> go ws []
        go [] _ls = mempty