{-# 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
   -- ^ The 'Column' from which the 'plainState_buffer'
   -- must be written.
 , plainState_bufferWidth     :: !Width
   -- ^ The 'Width' of the 'plainState_buffer' so far.
 , 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
	width = \case
	 PlainChunk_Ignored{} -> 0
	 PlainChunk_Word d -> width d
	 PlainChunk_Spaces s -> s
	nullWidth = \case
	 PlainChunk_Ignored{} -> True
	 PlainChunk_Word d -> nullWidth d
	 PlainChunk_Spaces s -> s == 0
instance From [SGR] d => From [SGR] (PlainChunk d) where
	from sgr = PlainChunk_Ignored (from 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 maxWidth | maxWidth < 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 maxWidth | maxWidth < newWidth ->
				overflow $ k ((spaces n <>), newState) fits fits
			 _ -> k ((spaces n <>), newState) fits overflow
instance (From (Word s) d, Semigroup d, Lengthable s) =>
         From (Word s) (Plain d) where
	from s = Plain $ \inh st@PlainState{..} k fits overflow ->
		let wordWidth = width s in
		if wordWidth <= 0
		then k (id,st) fits overflow
		else
			let newBufferWidth = plainState_bufferWidth + wordWidth in
			let newWidth = plainState_bufferStart + newBufferWidth in
			if plainInh_justify inh
			then
				let newState = st
					 { plainState_buffer =
						PlainChunk_Word (Word (from s)) :
						plainState_buffer
					 , plainState_bufferWidth = newBufferWidth
					 } in
				case plainInh_width inh of
				 Just maxWidth | maxWidth < 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 maxWidth | maxWidth < newWidth ->
					overflow $ k ((from s <>), newState) fits fits
				 _ -> k ((from s <>), newState) fits overflow
instance (From (Word s) d, Lengthable s, Spaceable d, Splitable s) =>
         From (Line s) (Plain d) where
	from =
		mconcat .
		List.intersperse breakspace .
		(from <$>) .
		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 = (\x -> flushLine <> x <> flushLine) $ Plain $ \inh ->
		unPlain p inh{plainInh_justify=True}
		where
		flushLine :: Plain d
		flushLine = 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 (From (Word String) d, Spaceable d) =>
         From String (Plain d) where
	from =
		mconcat .
		List.intersperse newline .
		(from <$>) .
		lines
instance (From (Word String) d, Spaceable d) =>
         IsString (Plain d) where
	fromString = from
-- Text
instance (From (Word Text) d, Spaceable d) =>
         From Text (Plain d) where
	from =
		mconcat .
		List.intersperse newline .
		(from <$>) .
		lines
instance (From (Word TL.Text) d, Spaceable d) =>
         From TL.Text (Plain d) where
	from =
		mconcat .
		List.intersperse newline .
		(from <$>) .
		lines
-- Char
instance (From (Word Char) d, Spaceable d) =>
         From Char (Plain d) where
	from ' '  = breakspace
	from '\n' = newline
	from c    = from (Word c)


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

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

countWords :: [PlainChunk d] -> Natural
countWords = go False 0
 where
	go inWord acc = \case
	 [] -> acc
	 PlainChunk_Word{}:xs ->
		if inWord
		then go inWord acc xs
		else go True (acc+1) xs
	 PlainChunk_Spaces s:xs
	  | s == 0    -> go inWord acc xs
	  | otherwise -> go False acc xs
	 PlainChunk_Ignored{}:xs -> go inWord acc xs

-- | @('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 bufferWidth (lineWidth,wordCount,line) = Line $
	if bufferWidth <= lineWidth
		-- The gathered line reached or overreached the bufferWidth,
		-- hence no padding id needed.
	|| wordCount <= 1
		-- The case bufferWidth <= lineWidth && wordCount == 1
		-- can happen if first word's length is < bufferWidth
		-- but second word's len is >= bufferWidth.
	then joinPlainLine line
	else
		-- Share the missing spaces as evenly as possible
		-- between the words of the line.
		padPlainLine line $ justifyPadding (bufferWidth-lineWidth) (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