{-# LANGUAGE TypeSynonymInstances, FlexibleInstances, TemplateHaskell, ConstraintKinds, FlexibleContexts #-}

module Text.MPretty.MonadPretty where

import Prelude hiding (id, (.))

import Data.PartialOrder
import qualified Data.List as L
import Data.Char
import Control.Category
import Control.Monad
import Control.Monad.Reader
import Control.Monad.State
import Control.Monad.Writer
import Data.Lens
import System.Console.ANSI
import Text.MPretty.StateSpace
import Util.ConsoleState
import Util.HasLens
import Util.Lens
import Util.List

----- "Primitives" -----

text :: (MonadPretty env out state m) => out -> m ()
text s = 
  let sL = pLength s
      nsL = countNonSpace s
  in do
    tell s
    columnL . view %= (+) sL
    ribbonL . view %= (+) nsL
    m <- look $ failureL . view
    when (m == Fail) $ do
      w <- look $ layoutWidthL . view
      rr <- look $ ribbonRatioL . view
      k <- access $ columnL . view
      r <- access $ ribbonL . view
      when (k > w) mzero
      when (fromIntegral r > fromIntegral w * rr) mzero
    where
      countNonSpace = pFoldl (\i c -> i + if isSpace c then 0 else 1) 0

string :: (MonadPretty env out state m) => String -> m ()
string = text . pString

space :: (MonadPretty env out state m) => Int -> m ()
space = string . flip replicate ' '

tryFlat :: (MonadPretty env out state m) => m a -> m a -> m a
tryFlat dFlat dBreak = do
  l <- look $ layoutL . view
  case l of
    Flat -> dFlat
    Break -> dBreak

hardLine :: (MonadPretty env out state m) => m ()
hardLine = do
  i <- look $ nestingL . view
  tell $ pString "\n"
  columnL . view ~= 0
  ribbonL . view ~= 0
  space i

flatFail :: (MonadPretty env out state m) => m a -> m a
flatFail = 
  local (modL (failureL . view) $ const Fail) 
  . local (modL (layoutL . view) $ const Flat)

nest :: (MonadPretty env out state m) => Int -> m a -> m a
nest i = local $ modL (nestingL . view) (i +)

group :: (MonadPretty env out state m) => m a -> m a
group aM = do
  l <- look $ layoutL . view
  case l of
    Flat -> aM
    Break -> msum
      [ flatFail aM
      , aM
      ]

align :: (MonadPretty env out state m) => m a -> m a
align aM = do
  i <- look $ nestingL . view
  k <- access $ columnL . view
  nest (k-i) aM

hang :: (MonadPretty env out state m) => Int -> m a -> m a
hang i = align . nest i

precedence :: (MonadPretty env out state m) => (Precedence,Precedence) -> m a -> m a
precedence = local . setL (precedenceL . view)

style :: (MonadPretty env out state m) => Style -> m a -> m a
style = local . setL (styleL . styleOptionsL . view)

buffering :: (MonadPretty env out state m) => Buffering -> m a -> m a
buffering = local . setL (bufferingL . styleOptionsL . view)

doConsole :: (MonadPretty env out state m) => Bool -> m a -> m a
doConsole = local . setL (doConsoleL . view)

layoutWidth :: (MonadPretty env out state m) => Int -> m a -> m a
layoutWidth = local . setL (layoutWidthL . view)

indentWidth :: (MonadPretty env out state m) => Int -> m a -> m a
indentWidth = local . setL (indentWidthL . styleOptionsL . view)

----- Helpers -----

buffer :: (MonadPretty env out state m) => m a -> m a
buffer = buffering Buffer

noBuffer :: (MonadPretty env out state m) => m a -> m a
noBuffer = buffering NoBuffer

console :: (MonadPretty env out state m) => m a -> m a
console = doConsole True

noConsole :: (MonadPretty env out state m) => m a -> m a
noConsole = doConsole False

closedPrecedence :: Int -> (Precedence,Precedence)
closedPrecedence i = (Precedence i NoD False,Precedence i NoD False)

getBuff :: (MonadPretty env out state m) => m out
getBuff = do
  b <- look $ bufferingL . styleOptionsL . view
  return $ case b of
    Buffer -> pString " "
    NoBuffer -> mempty

----- Style helpers -----

dropIndent :: (MonadPretty env out state m) => m () -> m ()
dropIndent d = do
  i <- look $ indentWidthL . styleOptionsL . view
  tryFlat (return ()) $ do
    hardLine
    space i
  align d

encloseSepPre :: (MonadPretty env out state m)
              => out -> out -> out -> Bool -> [m ()] -> m ()
encloseSepPre lbrac rbrac sep snug ds = 
  let lbracL = pLength lbrac
      sepL = pLength sep
  in do
    buff <- getBuff
    let f = foldr (.) id
          [ mapFirst $ \ d -> do
              punctuation $ text lbrac
              tryFlat (text buff) $ do 
                space $ sepL - lbracL
                text buff
              d
          , mapRest $ \ d -> do
              tryFlat (text buff) $ do
                hardLine
                space $ lbracL - sepL
              punctuation $ text sep
              text buff
              d
          , mapLast $ \ d -> do
              d
              if snug then text buff else tryFlat (text buff) hardLine
              punctuation $ text rbrac
          ]
    group . sequence_ . f $ map (precedence (closedPrecedence 0)  . align) ds

encloseSepPost :: (MonadPretty env out state m)
               => out -> out -> out -> [m ()] -> m ()
encloseSepPost lbrac rbrac sep ds =
  let lbracL = pLength lbrac
  in do
    buff <- getBuff
    let f = foldr (.) id $
          [ mapFirst $ \ d -> do
              punctuation $ text lbrac
              text buff
              d
          , mapRest $ \ d -> do
              tryFlat (return ()) $ do
                hardLine
                space lbracL
              text buff
              d
          , mapLeading $ \ d -> do
              d
              text buff
              punctuation $ text sep
          , mapLast $ \ d -> do
              d
              text buff
              punctuation $ text rbrac
          ]
    group . sequence_ . f $ map (precedence (closedPrecedence 0) . align) ds

encloseSepIndent :: (MonadPretty env out state m)
                 => out -> out -> out -> [m ()] -> m ()
encloseSepIndent lbrac rbrac sep ds = do
  buff <- getBuff
  i <- look $ indentWidthL . styleOptionsL . view
  let f = foldr (.) id $
        [ mapFirst $ \ d -> do
            punctuation $ text lbrac
            d
        , map $ \ d -> do
            tryFlat (text buff) $ do
              hardLine
              space i
            d
        , mapLeading $ \ d -> do
            d
            text buff
            punctuation $ text sep
        , mapLast $ \ d -> do
            d
            tryFlat (text buff) hardLine
            punctuation $ text rbrac
        ]
  group . sequence_ . f $ map (precedence (closedPrecedence 0) . align) ds

encloseSep :: (MonadPretty env out state m) 
           => out -> out -> out -> [m ()] -> m ()
encloseSep lbrac rbrac _ [] = punctuation $ text lbrac >> text rbrac
encloseSep lbrac rbrac sep ds = do
  s <- look $ styleL . styleOptionsL . view
  case s of
    PreAlignStyle -> encloseSepPre lbrac rbrac sep False ds
    PreSnugStyle -> encloseSepPre lbrac rbrac sep True ds
    PostStyle -> encloseSepPost lbrac rbrac sep ds
    IndentStyle -> encloseSepIndent lbrac rbrac sep ds

encloseSepDropIndent :: (MonadPretty env out state m)
                     => out -> out -> out -> [m ()] -> m ()
encloseSepDropIndent lbrac rbrac _ [] = punctuation $ text lbrac >> text rbrac
encloseSepDropIndent lbrac rbrac sep ds = do
  s <- look $ styleL . styleOptionsL . view
  case s of
    PreAlignStyle -> dropIndent $ encloseSepPre lbrac rbrac sep False ds
    PreSnugStyle -> dropIndent $ encloseSepPre lbrac rbrac sep True ds
    PostStyle -> dropIndent $ encloseSepPost lbrac rbrac sep ds
    IndentStyle -> encloseSepIndent lbrac rbrac sep ds

infixOp :: (MonadPretty env out state m) 
        => Direction -> Int -> Buffering -> m () -> m () -> m () -> m ()
infixOp d n b infixD leftD rightD = do
  s <- look $ styleL . styleOptionsL . view
  let buff = case b of
        Buffer -> pString " "
        NoBuffer -> mempty
  (pl,pr) <- look $ precedenceL . view
  let q = Precedence n d False
      ql = if d == LeftD then q else pbump q
      qr = if d == RightD then q else pbump q
      enclose = if lte pl q && lte pr q
        then id
        else group . parenthesize
  enclose $ do
    (pl',pr') <- look $ precedenceL . view
    precedence (pl',ql) leftD
    let preSep = do
          tryFlat (text buff) hardLine
          infixD
          text buff
        postSep = do
          text buff
          infixD
          tryFlat (text buff) hardLine
    case s of
      PreAlignStyle -> preSep
      PreSnugStyle -> preSep
      PostStyle -> postSep
      IndentStyle -> postSep
    precedence (qr,pr') rightD

hsep :: (MonadPretty env out state m) => [m ()] -> m ()
hsep ds = do
  buff <- getBuff
  foldr (>>) (return ()) $ L.intersperse (text buff) ds

vsep :: (MonadPretty env out state m) => [m ()] -> m ()
vsep ds = do
  buff <- getBuff
  foldr (>>) (return ()) $ L.intersperse (tryFlat (text buff) hardLine) ds

parenthesize :: (MonadPretty env out state m) => m () -> m ()
parenthesize d = do
  punctuation $ string "("
  precedence (closedPrecedence 0) $ align d
  punctuation $ string ")"

sexpListCons :: (MonadPretty env out state m) => [m ()] -> Maybe (m ()) -> m ()
sexpListCons ds dM = group $ parenthesize $ do
  buffer $ vsep $ ds
  case dM of
    Nothing -> return ()
    Just d -> do
      tryFlat (space 1) hardLine
      punctuation $ string ". "
      d

sexpList :: (MonadPretty env out state m) => [m ()] -> m ()
sexpList = flip sexpListCons Nothing

-- a pretty-printing mode suitable for providing a Show instance

showStyle :: (MonadPretty env out state m) => m a -> m a
showStyle = 
  layoutWidth 0 
  . local (setL (layoutL . view) Flat) 
  . noBuffer 
  . noConsole

----- ANSI Console helpers -----

emitConsoleStateCodes :: (MonadPretty env out state m) => m ()
emitConsoleStateCodes = do
  proceed <- look $ doConsoleL . view
  when proceed $ do
    cs <- look $ consoleStateL . view
    tell $ pString $ setConsoleStateCodes cs

localConsole :: (MonadPretty env out state m) => (ConsoleState -> ConsoleState) -> m a -> m a
localConsole f aM = do
  a <- local (modL (consoleStateL . view) f) $ do
    emitConsoleStateCodes
    aM
  emitConsoleStateCodes
  return a

intensity :: (MonadPretty env out state m) => ConsoleIntensity -> m a -> m a
intensity = localConsole . setL intensityML . Just

italicized :: (MonadPretty env out state m) => Bool -> m a -> m a
italicized = localConsole . setL italicizedML . Just

underlining :: (MonadPretty env out state m) => Underlining -> m a -> m a
underlining = localConsole . setL underliningML . Just

blinkSpeed :: (MonadPretty env out state m) => BlinkSpeed -> m a -> m a
blinkSpeed = localConsole . setL blinkSpeedML . Just

visible :: (MonadPretty env out state m) => Bool -> m a -> m a
visible = localConsole . setL visibleML . Just

swapFgBg :: (MonadPretty env out state m) => Bool -> m a -> m a
swapFgBg = localConsole . setL swapFgBgML . Just

gcolor :: (MonadPretty env out state m) => ConsoleLayer -> ColorIntensity -> Color -> m a -> m a
gcolor cl ci c = localConsole $ setL gcolorML $ Just (cl,ci,c)

color :: (MonadPretty env out state m) => ColorIntensity -> Color -> m a -> m a
color = gcolor Foreground

localStyle :: (MonadPretty env out state m) => Lens Palette ConsoleState -> m a -> m a
localStyle l aM = do
  c <- look $ l . paletteL . view
  localConsole (mappend c) aM

punctuation :: (MonadPretty env out state m) => m a -> m a
punctuation = localStyle punctuationColorL

literal :: (MonadPretty env out state m) => m a -> m a
literal = localStyle literalColorL

binder :: (MonadPretty env out state m) => m a -> m a
binder = localStyle binderColorL

keyword :: (MonadPretty env out state m) => m a -> m a
keyword = localStyle keywordColorL

classifier :: (MonadPretty env out state m) => m a -> m a
classifier = localStyle classifierColorL

----- Testing -----

styleVariants :: (MonadPretty env out state m) => m () -> m ()
styleVariants aM = do
  i <- look $ indentWidthL . styleOptionsL . view
  let configs =
        [ StyleOptions PreAlignStyle Buffer   i
        , StyleOptions PreAlignStyle NoBuffer i
        , StyleOptions PreSnugStyle  Buffer   i
        , StyleOptions PreSnugStyle  NoBuffer i
        , StyleOptions PostStyle     Buffer   i
        , StyleOptions PostStyle     NoBuffer i
        , StyleOptions IndentStyle   Buffer   i
        , StyleOptions IndentStyle   NoBuffer i
        ]
  forM_ configs $ \ o -> do
    hardLine
    string "##### "
    string $ show o
    string " #####"
    hardLine
    local (setL (styleOptionsL . view) o) aM
    hardLine