{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE OverloadedStrings #-}

import Control.Applicative
import Control.Monad
import Control.Monad.State
import Data.Attoparsec.Text hiding (take)
import Data.Char
  ( isAlpha,
    isDigit,
  )
import qualified Data.Text as T
import qualified Data.Text.IO as T
import Prelude hiding
  ( print,
    takeWhile,
  )

data Tokens = Wrd {unwrd :: T.Text} | Cmd T.Text | BSq [Tokens] | BCr [Tokens] | NL | Cmt T.Text | Nobs | IMath T.Text | Math T.Text | Verbatim Bool T.Text
  deriving (Show)

isCtrl x = not $ isAlpha x || isDigit x

-- tokens to strings
toStr (Verbatim star str) = do
  print' "\n\\begin{verbatim*}"
  print' str
  print' "\\end{verbatim*}"
toStr (Wrd w) = do
  if T.length w >= 1 && T.head w `elem` (".,;:!?)" :: [Char])
    then print' w
    else print w
  when (w == "(") suppress
  when (T.last w == '.') printnl
toStr (Cmd c) =
  if T.length c > 0 && isCtrl (T.head c) || T.length c == 1
    then print' $ "\\" `T.append` c
    else do
      printnl'
      when (c == "end") unindent
      print $ "\\" `T.append` c
      when (c == "begin") indent
toStr (BSq ts) = do
  print' "["
  suppress
  indent
  mapM_ toStr ts
  unindent
  print' "]"
toStr (BCr ts) = do
  print' "{"
  suppress
  indent
  mapM_ toStr ts
  unindent
  print' "}"
toStr NL = printnl' >> printnl
toStr (Cmt c) = do
  print "%"
  print' c
toStr Nobs = do
  print' "~"
  suppress
toStr (Math m) = do
  printnl
  print "\\["
  print' m
  println "\\]"
toStr (IMath m) = do
  print "\\("
  print' m
  print' "\\)"

-- Tokeniser
ws = inClass " \t\n"

pad = takeWhile ws

token = nl <|> pad *> (nobs <|> verbatim <|> verbatimStar <|> math <|> cmd <|> bsq <|> bcr <|> cmt <|> wrd)

nobs = string "~" *> pure Nobs

nl = string "\n" *> takeWhile (inClass " \t") *> string "\n" *> pure NL

cmd :: Parser Tokens
cmd = Cmd <$> (string "\\" *> ((unwrd <$> wrd) <|> pure ""))

bsq = BSq <$> (string "[" *> many' token) <* pad <* string "]"

bcr = BCr <$> (string "{" *> many' token) <* pad <* string "}"

wrd = Wrd <$> takeWhile1 (notInClass " \t\n[]{}\\~")

cmt = Cmt <$> (takeWhile1 (== '%') *> takeWhile (/= '\n'))

math = oldmath <|> newmath <|> oldimath <|> newimath

oldmath = Math <$> (string "$$" *> takeWhile (/= '$')) <* string "$$"

newmath = Math . T.pack <$> (string "\\[" *> manyTill' anyChar (string "\\]"))

oldimath = IMath <$> (string "$" *> takeWhile (/= '$')) <* string "$"

newimath = IMath . T.pack <$> (string "\\(" *> manyTill' anyChar (string "\\)"))

verbatim = Verbatim False . T.pack <$> (string "\\begin{verbatim}" *> manyTill' anyChar (string "\\end{verbatim}"))

verbatimStar = Verbatim True . T.pack <$> (string "\\begin{verbatim*}" *> manyTill' anyChar (string "\\end{verbatim*}"))

-- indented printer
type Beginning = Bool

type Indent = Int

type SuppressSpace = Bool

data Printer = Printer Beginning SuppressSpace Indent T.Text
  deriving (Show)

print :: T.Text -> State Printer ()
print str = do
  (Printer beg sup i s) <- get
  if sup
    then print' str
    else do
      let indent = T.pack . take (2 * i) $ repeat ' '
          str' = T.concat [s, if beg then indent else " ", str]
      put $ Printer False False i str'

-- force unindented
print' :: T.Text -> State Printer ()
print' str = do
  (Printer beg _ i s) <- get
  let indent = T.pack . take (2 * i) $ repeat ' '
      str' = T.concat [s, if beg then indent else "", str]
  put $ Printer False False i str'

printnl :: State Printer ()
printnl = do
  (Printer beg sup i str) <- get
  when (not sup) $ put $ Printer True False i $ str `T.append` "\n"

printnl' :: State Printer ()
printnl' = do
  (Printer beg sup i str) <- get
  when (not beg && not sup) $ put $ Printer True False i $ str `T.append` "\n"

println :: T.Text -> State Printer ()
println str = do
  print str
  printnl

println' :: T.Text -> State Printer ()
println' str = do
  print' str
  printnl

indent :: State Printer ()
indent = do
  (Printer b sup i s) <- get
  put $ Printer b sup (i + 1) s

unindent :: State Printer ()
unindent = do
  (Printer b sup i s) <- get
  put $ Printer b sup (i - 1) s

suppress :: State Printer ()
suppress = do
  (Printer b _ i s) <- get
  put $ Printer b True i s

-- entry

main = do
  str <- T.getContents
  let Right toks = parseOnly (many' token) str
      Printer _ _ _ str' =
        flip execState (Printer True False 0 "") $ mapM_ toStr toks
  T.putStr str'