-- | The old CTPL0 virtual machine. This one is outdated, but will still be supported for a while.

module Text.CTPL0 (Exec(..), CTPL0(..), CTPL0State(..), BufferState(..), RegisterState(..), InfoState(..), unetx, endOfInstr, singleInstr, procInstrs, evalCTPL0', evalCTPL0) where

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

-- | A character buffer. Represented as a triplet of left behind chars, the current char and the pending chars.
data BufferState = BufferState {
  leftBehind :: String,     -- ^ Already seen. String reversed!
  thisChar :: Char,         -- ^ The current char.
  rightPending :: String    -- ^ The pending chars.
  }

-- | A record for register and stack state.
data RegisterState = RegisterState {
  ax :: Integer,   -- ^ Accumulator register (AX)
  mk :: [String],  -- ^ Tape stack (MK)
  rk :: [Int],     -- ^ Return address stack (RK)
  ck :: [Int],     -- ^ Number stack (CK). The top element may be used like a register.
  cp :: Bool       -- ^ If set, the next instruction will use the top element of CK instead of AX. (Say: ``SX is CK top'' in contrast to ``SX is AX'')
  }

-- | A statistics record (might be used by a profiler some time in future...)
data InfoState = InfoState {
  instrStats :: AVL (Char, Int) -- ^ Statistics on how often each instruction has been executed.
  }

-- | The overall state record.
data CTPL0State = CTPL0State {
  bufferState :: BufferState,     -- ^ Tape buffer.
  programState :: BufferState,    -- ^ Program buffer.
  registerState :: RegisterState, -- ^ Register state record.
  infoState :: InfoState          -- ^ Statistics record.
  }

-- | Monad displaying success or failure.
data Exec a = Succ a    -- ^ Execution succeeded :)
            | Expired   -- ^ Nope. Time has expired. Program took too long to finish. You might want to increase time limit.
            | ConfViol  -- ^ Nope. Confidence violation. This may have several reasons, e.g. popping from an empty stack, jumping out of program bounds, ...
            | SynViol   -- ^ Nope. Syntax violation. I encountered an instruction (or condition) I do not understand.
            deriving Show

instance Monad Exec where
  return = Succ
  (Succ a) >>= f = f a
  Expired >>= f = Expired
  ConfViol >>= f = ConfViol
  SynViol >>= f = SynViol

instance Applicative Exec where
  pure = return
  (<*>) = ap

instance Functor Exec where
  fmap = liftM

-- | The VM's execution monad. Behaves like a 'StateT' carrying a 'CTPL0State' wrapped around the 'Exec' monad. Also responsible for time consumption and passing.
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

instance Applicative CTPL0 where
  pure = return
  (<*>) = ap

instance Functor CTPL0 where
  fmap = liftM

-- | Gets the carried 'CTPL0State' and runs a function on it.
getState :: (CTPL0State -> a) -> CTPL0 a
getState f = CTPL0 $ \i k -> Succ (f k, k, i)

-- | Runs a function on the carried 'CTPL0State'.
modState :: (CTPL0State -> CTPL0State) -> CTPL0 ()
modState f = CTPL0 $ \i k -> Succ ((), f k, i)

-- | Consume virtual time. Raise 'Expired' if limit is reached.
consumeTime :: CTPL0 ()
consumeTime = CTPL0 $ \i k -> if i >= 1 then Succ ((), k, i-1) else Expired

-- | Raise a 'ConfViol'.
confViol :: CTPL0 a
confViol = CTPL0 $ \_ _ -> ConfViol

-- | Raise a 'SynViol'.
synViol :: CTPL0 a
synViol = CTPL0 $ \_ _ -> SynViol

-- | Modify the tape buffer`s state by running a function on it.
modBufferState :: (BufferState -> BufferState) -> CTPL0 ()
modBufferState f = modState $ \s -> s{bufferState = f $ bufferState s}

-- | Modify the program buffer`s state by running a function on it.
modProgramState :: (BufferState -> BufferState) -> CTPL0 ()
modProgramState f = modState $ \s -> s{programState = f $ programState s}

-- | Modify the register state record by running a function on it.
modRegisterState :: (RegisterState -> RegisterState) -> CTPL0 ()
modRegisterState f = modState $ \s -> s{registerState = f $ registerState s}

-- | Walk in the buffer. A positive number specifies walking to the right, a negative one to the left.
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)

-- | Fetch the next instruction.
getInstr :: CTPL0 Char
getInstr = do
  k <- getState $ thisChar . programState
  modProgramState $ walkBuffer 1
  return k

-- | Have we reached the end of the program tape?
endOfInstr :: CTPL0 Bool
endOfInstr = getState $ null . rightPending . programState

-- | Fetch numeric argument (as many digits as we can get)
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

-- | Fetch string argument (delimited by '$')
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

-- | Get position in program buffer.
getIP :: CTPL0 Int
getIP = getState $ length . leftBehind . programState

-- | Get position in tape buffer.
getCP :: CTPL0 Int
getCP = getState $ length . leftBehind . bufferState

-- | Are we able to walk that far in the program buffer?
canRelJump :: Int -> CTPL0 Bool
canRelJump 0 = return True
canRelJump i
  | i < 0 = getState $ (>= -i) . length . leftBehind . programState
  | i > 0 = getState $ (>= i) . length . rightPending . programState

-- | Are we able to walk that far in the tape buffer?
canRelWalk :: Int -> CTPL0 Bool
canRelWalk 0 = return True
canRelWalk i
  | i < 0 = getState $ (>= -i) . length . leftBehind . bufferState
  | i > 0 = getState $ (>= i) . length . rightPending . bufferState

-- | Run an action (first arg) iff the test (second arg) succeeds. Raise 'ConfViol' otherwise.
provided :: CTPL0 a -> CTPL0 Bool -> CTPL0 a
provided act test = do
  b <- test
  if b
     then act
     else confViol

-- | SX is AX by default, but CK top after `C'
sx :: RegisterState -> Integer
sx r | cp r = fromIntegral $ head $ ck r
sx r        = ax r

-- | Set SX (AX or CK top) value.
setSX :: Integer -> RegisterState -> RegisterState
setSX i r | cp r = r{ck=fromIntegral i : tail (ck r)}
setSX i r        = r{ax=i}

-- | Run the next instruction in program.
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

-- | Run the entire program.
procInstrs :: CTPL0 ()
procInstrs = singleInstr `asLongAs` liftM not endOfInstr
  where asLongAs act test = do
          b <- test
          when b $ act >> asLongAs act test

-- | A handy wrapper around 'procInstrs'. Arguments: program, tape, time limit. Results: tape, leftover time, AX, CK top, instruction stats.
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

-- | Another handy wrapper around 'procInstr'. Less clumsy than 'evalCTPL0'', but provides less information. Arguments: program, tape, time limit. Results: tape only.
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

-- | Remove a trailing ETX (if there is one).
unetx :: String -> String
unetx [] = []
unetx s
  | s == [chr 3] = []
  | last s == chr 3 = init s
  | otherwise = s

-- | Append right before the trailing ETX. If there is no ETX, string will be appended at the end.
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