module Text.CTPL0 where

import Control.Monad
import Data.Char
import Data.Chatty.AVL
import Data.List

data BufferState = BufferState { leftBehind :: String, thisChar :: Char, rightPending :: String }
data RegisterState = RegisterState { ax :: Integer, mk :: [String], rk :: [Int], ck :: [Int], cp :: Bool }
data InfoState = InfoState { instrStats :: AVL (Char, Int) }
data CTPL0State = CTPL0State { bufferState :: BufferState, programState :: BufferState, registerState :: RegisterState, infoState :: InfoState }
data Exec a = Succ a | Expired | ConfViol | SynViol deriving Show
newtype CTPL0 a = CTPL0 { runCTPL0 :: Int -> CTPL0State -> Exec (a, CTPL0State, Int) }

instance Monad CTPL0 where
  return a = CTPL0 $ \i k -> Succ (a, k, i)
  m >>= f = CTPL0 $ \i k ->
    case runCTPL0 m i k of
      Succ (a, k', i') -> runCTPL0 (f a) i' k'
      Expired -> Expired
      ConfViol -> ConfViol
      SynViol -> SynViol

getState :: (CTPL0State -> a) -> CTPL0 a
getState f = CTPL0 $ \i k -> Succ (f k, k, i)

modState :: (CTPL0State -> CTPL0State) -> CTPL0 ()
modState f = CTPL0 $ \i k -> Succ ((), f k, i)

consumeTime :: CTPL0 ()
consumeTime = CTPL0 $ \i k -> if i >= 1 then Succ ((), k, i-1) else Expired

confViol :: CTPL0 a
confViol = CTPL0 $ \_ _ -> ConfViol

synViol :: CTPL0 a
synViol = CTPL0 $ \_ _ -> SynViol

modBufferState :: (BufferState -> BufferState) -> CTPL0 ()
modBufferState f = modState $ \s -> s{bufferState = f $ bufferState s}

modProgramState :: (BufferState -> BufferState) -> CTPL0 ()
modProgramState f = modState $ \s -> s{programState = f $ programState s}

modRegisterState :: (RegisterState -> RegisterState) -> CTPL0 ()
modRegisterState f = modState $ \s -> s{registerState = f $ registerState s}

walkBuffer :: Int -> BufferState -> BufferState
walkBuffer 0 s = s
walkBuffer i s
  | i < 0 = BufferState (drop (-i) $ leftBehind s) (head $ drop (-i-1) $ leftBehind s) (reverse (take (-i-1) $ leftBehind s) ++ [thisChar s] ++ rightPending s)
  | i > 0 = BufferState (reverse (take (i-1) (rightPending s)) ++ [thisChar s] ++ leftBehind s) (head $ drop (i-1) $ rightPending s) (drop i $ rightPending s)

getInstr :: CTPL0 Char
getInstr = do
  k <- getState $ thisChar . programState
  modProgramState $ walkBuffer 1
  return k

endOfInstr :: CTPL0 Bool
endOfInstr = getState $ null . rightPending . programState

instrNumArg :: CTPL0 Int
instrNumArg = do
  ks <- getState $ \s -> takeWhile isDigit (thisChar (programState s) : rightPending (programState s))
  when (null ks) synViol
  modProgramState $ walkBuffer $ length ks
  return $ read ks

instrDelimArg :: CTPL0 String
instrDelimArg = do
  ks <- getState $ \s -> takeWhile (/='$') (thisChar (programState s) : rightPending (programState s))
  modProgramState $ walkBuffer $ length ks
  k' <- getState $ (=='$') . thisChar . programState
  unless k' synViol
  modProgramState $ walkBuffer 1
  return ks

getIP :: CTPL0 Int
getIP = getState $ length . leftBehind . programState

getCP :: CTPL0 Int
getCP = getState $ length . leftBehind . bufferState

canRelJump :: Int -> CTPL0 Bool
canRelJump 0 = return True
canRelJump i
  | i < 0 = getState $ (>= -i) . length . leftBehind . programState
  | i > 0 = getState $ (>= i) . length . rightPending . programState

canRelWalk :: Int -> CTPL0 Bool
canRelWalk 0 = return True
canRelWalk i
  | i < 0 = getState $ (>= -i) . length . leftBehind . bufferState
  | i > 0 = getState $ (>= i) . length . rightPending . bufferState

provided :: CTPL0 a -> CTPL0 Bool -> CTPL0 a
provided act test = do
  b <- test
  if b
     then act
     else confViol

-- SX points to AX by default, but to CK(0) after 'C'
sx :: RegisterState -> Integer
sx r | cp r = fromIntegral $ head $ ck r
sx r        = ax r

setSX :: Integer -> RegisterState -> RegisterState
setSX i r | cp r = r{ck=fromIntegral i : tail (ck r)}
setSX i r        = r{ax=i}

singleInstr :: CTPL0 ()
singleInstr = do
  i <- getInstr
  consumeTime
  f <- getState $ instrStats . infoState
  let f' = case avlLookup i f of
        Nothing -> avlInsert (i,1) f
        Just j -> avlInsert (i,j+1) f
  modState $ \s -> s{infoState=InfoState f'}
  case i of
    -- Walk left
    '<' -> modBufferState (walkBuffer (-1))
           `provided` getState (not . null . leftBehind . bufferState)
    -- Walk right
    '>' -> modBufferState (walkBuffer 1)
           `provided` getState (not . null . rightPending . bufferState)
    -- Inc AX (CK(0))
    '+' -> do
      num <- liftM fromIntegral instrNumArg
      modRegisterState $ \s -> setSX (sx s + num) s
    -- Dec AX (CK(0))
    '-' -> do
      num <- liftM fromIntegral instrNumArg
      modRegisterState $ \s -> setSX (sx s - num) s
    -- Insert char, go after
    'i' -> do
      ch <- getInstr `provided` liftM not endOfInstr
      modBufferState $ \s -> s{leftBehind=ch : leftBehind s}
    -- Replace char
    'r' -> do
      ch <- getInstr `provided` liftM not endOfInstr
      modBufferState $ \s -> s{thisChar=ch}
    -- Delete char
    'x' -> modBufferState (\s -> s{thisChar=head $ rightPending s, rightPending=tail $ rightPending s})
           `provided` getState (not . null . rightPending . bufferState)
    -- Insert chars delimited by $, go after
    'I' -> do
      cs <- instrDelimArg
      modBufferState $ \s -> s{leftBehind = reverse cs ++ leftBehind s}
    -- Append char at the end, don't walk
    'a' -> do
      ch <- getInstr `provided` liftM not endOfInstr
      modBufferState $ \s -> s{rightPending=appendBeforeETX (rightPending s) [ch]}
    -- Append chars delimited by $, don't walk
    'A' -> do
      cs <- instrDelimArg
      modBufferState $ \s -> s{rightPending=appendBeforeETX (rightPending s) cs}
    -- Push [CP] to MK
    'y' -> do
      ch <- getState $ thisChar . bufferState
      modRegisterState $ \s -> s{mk=[ch]:mk s}
    -- Append [CP] to MK(0)
    'Y' -> do
      ch <- getState $ thisChar . bufferState
      modRegisterState (\s -> s{mk=(appendBeforeETX (head $ mk s) [ch]):tail (mk s)})
       `provided` getState (not . null . mk . registerState)
    -- Pop MK(0), discard
    'p' -> modRegisterState (\s -> s{mk=tail $ mk s})
           `provided` getState (not . null . mk . registerState)
    -- Peek MK(0), insert, go after
    'P' -> do
      cs <- getState (head . mk . registerState)
            `provided` getState (not . null . mk . registerState)
      modBufferState $ \s -> s{leftBehind = reverse (unetx cs) ++ leftBehind s}
    -- Set IP = AX (CK(0))
    'j' -> do
      ax <- getState $ sx . registerState
      b <- singleCond
      when b $ do
        ip <- getIP
        let rel = fromIntegral ax - ip
        modProgramState (walkBuffer rel) `provided` canRelJump rel
    -- Set IP += AX (CK(0))
    'J' -> do
      ax <- getState $ sx . registerState
      b <- singleCond
      when b $ do
        let rel = fromIntegral ax
        modProgramState (walkBuffer rel) `provided` canRelJump rel
    -- Set IP = AX (CK(0)), push IP onto RK
    'c' -> do
      ax <- getState $ sx . registerState
      b <- singleCond
      when b $ do
        ip <- getIP
        let rel = fromIntegral ax - ip
        modProgramState (walkBuffer rel) `provided` canRelJump rel
        modRegisterState $ \s -> s{rk=ip:rk s}
    -- Return to RK(0), pop RK
    'f' -> do
      r0 <- getState (head . rk . registerState)
            `provided` getState (not . null . rk . registerState)
      ip <- getIP
      let rel = r0 - ip
      modProgramState (walkBuffer rel) `provided` canRelJump rel
      modRegisterState $ \s -> s{rk=tail $ rk s}
    -- Set AX (CK(0)) = 0
    '0' -> modRegisterState $ setSX 0
    -- Set AX (CK(0)) = CP
    'Q' -> do
      cp <- getCP
      modRegisterState $ setSX $ fromIntegral cp
    -- Set CP = AX (CK(0))
    'm' -> do
      ax <- getState $ fromIntegral . sx . registerState
      cp <- getCP
      let rel = ax - cp
      modBufferState (walkBuffer rel) `provided` canRelWalk rel
    -- Select CK(0) instead of AX for next operation
    'C' -> modRegisterState (\s -> s{cp=True})
           `provided` getState (not . null . ck . registerState)
    -- Load ord[CP] into AX (CK(0))
    'l' -> do
      ch <- getState $ ord . thisChar . bufferState
      modRegisterState $ setSX $ fromIntegral ch
    -- Save ascii(AX) (CK(0)) to [CP]
    's' -> do
      ax <- getState $ fromIntegral . sx . registerState
      modBufferState $ \s -> s{thisChar=chr ax}
    -- Push AX onto CK (or duplicate CK0, if SX->CK0)
    'd' -> do
      ax <- getState $ fromIntegral . sx . registerState
      modRegisterState $ \s -> s{ck=ax:ck s}
    -- Pop AX from CK
    'D' -> do
      ax' <- getState (fromIntegral . head . ck . registerState)
             `provided` getState (not . null . ck . registerState)
      modRegisterState $ \s -> s{ax=ax',ck=tail (ck s)}
    -- Pop CK, discard
    'k' -> modRegisterState (\s -> s{ck=tail (ck s)})
           `provided` getState (not . null . ck . registerState)
    -- Catch others
    o -> synViol
  unless (i=='C') $ modRegisterState $ \s -> s{cp=False}

singleCond :: CTPL0 Bool
singleCond = do
  i <- getInstr `provided` liftM not endOfInstr
  case i of
    -- Is Uppercase?
    'U' -> getState $ isUpper . thisChar . bufferState
    -- Is Lowercase?
    'L' -> getState $ isLower . thisChar . bufferState
    -- AX (CK(0)) = 0 ?
    'z' -> getState $ (==0) . sx . registerState
    -- Always true
    't' -> return True
    -- Is Digit?
    'N' -> getState $ isDigit . thisChar . bufferState
    -- Is End of Buffer?
    'e' -> getState $ null . rightPending . bufferState
    -- Negation
    '!' -> liftM not singleCond
    -- Disjunction
    '|' -> liftM2 (||) singleCond singleCond
    -- Conjunction
    '&' -> liftM2 (&&) singleCond singleCond
    -- Given char equals [CP]
    'q' -> do
      ch <- getInstr `provided` liftM not endOfInstr
      getState $ (==ch) . thisChar . bufferState
    -- CP < AX (CK(0))
    'l' -> do
      cp <- getCP
      getState $ (cp <) . fromIntegral . sx . registerState
    -- CP > AX (CK(0))
    'g' -> do
      cp <- getCP
      getState $ (cp >) . fromIntegral . sx . registerState
    -- CP = AX (CK(0))
    'E' -> do
      cp <- getCP
      getState $ (cp ==) . fromIntegral . sx . registerState
    -- If SX->AX then make SX->CK(0), otherwise make SX->AX
    'C' -> do
      sxp <- getState $ cp . registerState
      if sxp
         then modRegisterState (\s -> s{cp=False})
         else modRegisterState (\s -> s{cp=True})
              `provided` getState (not . null . ck . registerState)
      singleCond
    -- AX = CK(0)?
    '=' -> liftM2 (==)
           (getState $ ax . registerState)
           (getState $ fromIntegral . head . ck . registerState)
           `provided` getState (not . null . ck . registerState)
    -- AX < CK(0)?
    '<' -> liftM2 (<)
           (getState $ ax . registerState)
           (getState $ fromIntegral . head . ck . registerState)
           `provided` getState (not . null . ck . registerState)
    -- AX > CK(0)?
    '>' -> liftM2 (>)
           (getState $ ax . registerState)
           (getState $ fromIntegral . head . ck . registerState)
           `provided` getState (not . null . ck . registerState)
    -- Pop CK, discard, then continue evaluation
    'k' -> do
      modRegisterState (\s -> s{ck=tail $ ck s})
       `provided` getState (not . null . ck . registerState)
      singleCond
    -- Catch others
    o -> synViol

procInstrs :: CTPL0 ()
procInstrs = singleInstr `asLongAs` liftM not endOfInstr
  where asLongAs act test = do
          b <- test
          when b $ act >> asLongAs act test

evalCTPL0' :: String -> String -> Int -> Exec (String, Int, Integer, Int, [] (Char, Int))
evalCTPL0' program buffer limit =
  let state0 = CTPL0State buffer0 program0 register0 info0
      buffer0
        | null buffer = BufferState [] (chr 3) []
        | otherwise = BufferState [] (head buffer) (tail buffer ++ [chr 3])
      program0
        | null program = BufferState [] (chr 3) []
        | otherwise = BufferState [] (head program) (tail program ++ [chr 3])
      register0 = RegisterState 0 [] [length program] [0] False
      info0 = InfoState EmptyAVL
      imprf avl = sortBy (\b a -> snd a `compare` snd b) $ avlInorder avl
  in case runCTPL0 procInstrs limit state0 of
    Succ (_, CTPL0State b p r f, i) -> Succ (unetx (reverse (leftBehind b) ++ [thisChar b] ++ rightPending b), i, ax r, head $ ck r, imprf $ instrStats f)
    ConfViol -> ConfViol
    SynViol -> SynViol
    Expired -> Expired

evalCTPL0 :: String -> String -> Int -> Exec String
evalCTPL0 program buffer limit =
  case evalCTPL0' program buffer limit of
    Succ (s,_,_,_,_) -> Succ s
    ConfViol -> ConfViol
    SynViol -> SynViol
    Expired -> Expired

unetx :: String -> String
unetx [] = []
unetx s
  | s == [chr 3] = []
  | last s == chr 3 = init s
  | otherwise = s

appendBeforeETX :: String -> String -> String
appendBeforeETX [] t = t
appendBeforeETX s t
  | s == [chr 3] = t++[chr 3]
  | last s == chr 3 = init s ++ t ++ [chr 3]
  | otherwise = s ++ t