module DSL where

import Data.Char (isDigit, isSpace)
import Data.Foldable (foldlM)

import Parsers

------------------------------------------------------------------------------
-- Data types / instances
------------------------------------------------------------------------------

data StackData
  = StackInt Integer
  | StackBool Bool

instance Show StackData where
  show (StackInt x)      = show x
  show (StackBool True)  = "true"
  show (StackBool False) = "false"

type TypeCheck = StackData -> Bool

type Stack = [StackData]

data Intrinsic
  -- core stack operations
  = I_DUMP
  | I_DROP
  | I_SWAP
  | I_DUP
  | I_OVER
  | I_ROT
  -- core arithmetic operations
  | I_PLUS
  | I_MINUS
  | I_TIMES
  | I_DIVMOD
  -- core boolean operations
  | I_NOT
  | I_AND
  | I_OR
  | I_XOR
    deriving (Show, Eq)

data Operation
  = OpPushData StackData
  | OpIntrinsic Intrinsic
    deriving (Show)

data StackModifier = StackModifier { smName :: String
                                   , smTypes :: [TypeCheck]
                                   , smFunc :: Stack -> IO Stack
                                   }

data Block
  = BLinear [Operation]
  | BIf Block Block 
    deriving (Show)

type Program = [Block]

data Machine = Machine { ok :: Bool
                       , stack :: Stack
                       }

data TokenData
  = T_WHITESPACE
  | T_IF
  | T_DO
  | T_END
  | T_BOOL_LITERAL Bool
  | T_INT_LITERAL Integer
  | T_INTRINSIC Intrinsic
    deriving (Show, Eq)

data DSLToken = DSLToken { tStr :: String -- original text
                         , tData :: TokenData -- actual data
                         } deriving (Show)

type DSLLexer  = Parser Char DSLToken
type DSLParser = Parser DSLToken

------------------------------------------------------------------------------
-- Lexing
------------------------------------------------------------------------------

buildDSLLexer :: (Parser Char [Char]) -> ([Char] -> TokenData) -> DSLLexer
buildDSLLexer p f = do
  str <- p
  return DSLToken { tStr=str
                  , tData=f str
                  }

keywordL :: String -> TokenData -> DSLLexer
keywordL s d = buildDSLLexer (list s) (const d)

intrinsicL :: String -> Intrinsic -> DSLLexer
intrinsicL s i = keywordL s $ T_INTRINSIC i

wsL :: DSLLexer
wsL = buildDSLLexer (mult1 $ satisfy isSpace) (const T_WHITESPACE)

boolLiteralL :: DSLLexer
boolLiteralL = t `alt` f
  where
    t = buildDSLLexer (list "true") (const $ T_BOOL_LITERAL True)
    f = buildDSLLexer (list "false") (const $ T_BOOL_LITERAL False)

intLiteralL :: DSLLexer
intLiteralL = buildDSLLexer go (T_INT_LITERAL . read)
  where
    go = do
      sign   <- optional $ token '-'
      digits <- mult1 $ satisfy isDigit
      result $ maybe digits (:digits) sign

mainLexer :: Parser Char [DSLToken]
mainLexer = phrase $ mult1 $ firstOf subLexers
  where
    subLexers = [wsL]
             ++ keywords
             ++ literals
             ++ intrinsics
    keywords = map (uncurry keywordL) [ ("IF",  T_IF)
                                      , ("DO",  T_DO)
                                      , ("END", T_END)
                                      ]
    literals = [ boolLiteralL
               , intLiteralL
               ]
    intrinsics = map (uncurry intrinsicL) [ (".",    I_DUMP)
                                          , ("DROP", I_DROP)
                                          , ("SWAP", I_SWAP)
                                          , ("DUP",  I_DUP)
                                          , ("OVER", I_OVER)
                                          , ("ROT",  I_ROT)
                                          , ("+",    I_PLUS)
                                          , ("-",    I_MINUS)
                                          , ("*",    I_TIMES)
                                          , ("/%",   I_DIVMOD)
                                          , ("!",    I_NOT)
                                          , ("&&",   I_AND)
                                          , ("||",   I_OR)
                                          , ("^",    I_XOR)
                                          ]

stripWhitespace :: [DSLToken] -> [DSLToken]
stripWhitespace = filter $ not . (==T_WHITESPACE) . tData

------------------------------------------------------------------------------
-- Parsing
------------------------------------------------------------------------------

wsP :: DSLParser ()
wsP = do
  t <- anyToken
  case tData t of
    T_WHITESPACE -> result ()
    _            -> flunk

dataLiteralP :: DSLParser StackData
dataLiteralP = do
  t <- anyToken
  case tData t of
    T_INT_LITERAL x  -> result $ StackInt x
    T_BOOL_LITERAL x -> result $ StackBool x
    _                -> flunk

pushDataP :: DSLParser Operation
pushDataP = OpPushData <$> dataLiteralP

intrinsicP :: DSLParser Operation
intrinsicP = do
  t <- anyToken
  case tData t of
    T_INTRINSIC i -> result $ OpIntrinsic i
    _             -> flunk

operationP :: DSLParser Operation
operationP = firstOf [ pushDataP
                     , intrinsicP
                     ]

linearP :: DSLParser Block
linearP = BLinear <$> mult1 operationP

ifP :: DSLParser Block
ifP = do
  _ <- satisfy $ (==T_IF) . tData
  c <- recover (BLinear []) blockP <* satisfy ((==T_DO) . tData)
  b <- blockP <* satisfy ((==T_END) . tData)
  return $ BIf c b

blockP :: DSLParser Block
blockP = firstOf [ ifP
                 , linearP
                 ]

programP :: DSLParser Program
programP = phrase $ mult1 blockP

stringToProgram :: String -> Maybe Program
stringToProgram str = do
  (_, tokens)  <- parse mainLexer str
  (_, program) <- parse programP $ stripWhitespace tokens
  return program

------------------------------------------------------------------------------
-- Type checking
------------------------------------------------------------------------------

tAny :: TypeCheck
tAny = const True

tBool :: TypeCheck
tBool (StackBool _) = True
tBool _             = False

tInt :: TypeCheck
tInt (StackInt _) = True
tInt _            = False

runChecks :: [TypeCheck] -> Stack -> Maybe String
runChecks fs s
  | length fs > length s        = Just "stack underflow"
  | not (and $ zipWith id fs s) = Just "type mis-match"
  | otherwise                   = Nothing

------------------------------------------------------------------------------
-- Intrinsics
------------------------------------------------------------------------------

dump :: StackModifier
dump = StackModifier { smName="DUMP", smTypes=ts, smFunc=f }
  where
    ts = [tAny]
    f (x:xs) = putStrLn (show x) >> return xs
    f _      = unreachable

drop' :: StackModifier
drop' = StackModifier { smName="DROP", smTypes=ts, smFunc=f }
  where
    ts = [tAny]
    f (_:xs) = return xs
    f _      = unreachable

swap :: StackModifier
swap = StackModifier { smName="SWAP", smTypes=ts, smFunc=f }
  where
    ts = [tAny, tAny]
    f (x:y:xs) = return $ y:x:xs
    f _        = unreachable

dup :: StackModifier
dup = StackModifier { smName="DUP", smTypes=ts, smFunc=f }
  where
    ts = [tAny]
    f (x:xs) = return $ x:x:xs
    f _      = unreachable

over :: StackModifier
over = StackModifier { smName="OVER", smTypes=ts, smFunc=f }
  where
    ts = [tAny, tAny]
    f (x:y:xs) = return $ y:x:y:xs
    f _        = unreachable

rot :: StackModifier
rot = StackModifier { smName="ROT", smTypes=ts, smFunc=f }
  where
    ts = [tAny, tAny, tAny]
    f (x:y:z:xs) = return $ y:z:x:xs
    f _          = unreachable

binArithmetic :: String -> (Integer -> Integer -> [Integer]) -> StackModifier
binArithmetic name f = StackModifier { smName=name, smTypes=ts, smFunc=f' }
  where
    ts = [tInt, tInt]
    f' (StackInt x:StackInt y:xs) = return $ (StackInt <$> f x y) ++ xs
    f' _                          = unreachable

plus :: StackModifier
plus = binArithmetic "PLUS" $ \ x y -> [x + y]

minus :: StackModifier
minus = binArithmetic "MINUS" $ \ x y -> [x - y]

times :: StackModifier
times = binArithmetic "TIMES" $ \ x y -> [x * y]

divMod' :: StackModifier
divMod' = binArithmetic "DIVMOD" $ \ x y -> [mod x y, div x y]

not' :: StackModifier
not' = StackModifier { smName="NOT", smTypes=ts, smFunc=f }
  where
    ts = [tBool]
    f (StackBool x:xs) = return $ StackBool (not x):xs
    f _                = unreachable

binBoolean :: String -> (Bool -> Bool -> [Bool]) -> StackModifier
binBoolean name f = StackModifier { smName=name, smTypes=ts, smFunc=f' }
  where
    ts = [tBool, tBool]
    f' (StackBool x:StackBool y:xs) = return $ (StackBool <$> f x y) ++ xs
    f' _                            = unreachable

and' :: StackModifier
and' = binBoolean "AND" $ \ x y -> [x && y]

or' :: StackModifier
or' = binBoolean "OR" $ \ x y -> [x || y]

xor :: StackModifier
xor = binBoolean "XOR" $ \ x y -> [x /= y]

------------------------------------------------------------------------------
-- Core operations
------------------------------------------------------------------------------

-- "halt and catch fire"
hcf :: Machine -> String -> IO Machine
hcf m msg = putStrLn msg >> return m{ ok=False }

unreachable :: a
unreachable = error "this branch should be unreachable"

pushData :: Machine -> StackData -> Machine
pushData m@Machine{ stack=xs } x = m{ stack=x:xs }

runModifier :: StackModifier -> Stack -> IO (Either String Stack)
runModifier sm s = case runChecks (smTypes sm) s of
  Just err -> return $ Left $ smName sm ++ ": " ++ err
  Nothing  -> Right <$> smFunc sm s

applyIntrinsic :: Intrinsic -> Machine -> IO Machine
applyIntrinsic i m = do
  res <- runModifier sm (stack m)
  case res of
    Left err -> hcf m err
    Right s' -> return m{ stack=s' }
  where
    sm = case i of
      I_DUMP   -> dump
      I_DROP   -> drop'
      I_SWAP   -> swap
      I_DUP    -> dup
      I_OVER   -> over
      I_ROT    -> rot
      I_PLUS   -> plus
      I_MINUS  -> minus
      I_TIMES  -> times
      I_DIVMOD -> divMod'
      I_NOT    -> not'
      I_AND    -> and'
      I_OR     -> or'
      I_XOR    -> xor

------------------------------------------------------------------------------
-- Interpretation
------------------------------------------------------------------------------

newMachine :: Machine
newMachine = Machine { ok=True, stack=[] }

applyOperation :: Machine -> Operation -> IO Machine
-- take no action if a previous step failed
applyOperation m@Machine{ ok=False } _   = return m
applyOperation m (OpPushData x)          = return $ pushData m x
applyOperation m (OpIntrinsic i)         = applyIntrinsic i m

applyLinear :: Machine -> [Operation] -> IO Machine
applyLinear = foldlM applyOperation

applyIf :: Block -> Block -> Machine -> IO Machine
applyIf _ _ m@Machine{ ok=False } = return m
applyIf c b m                     = do
  m' <- evalBlocks [c] m
  case m' of
    Machine{ ok=False }                 -> return m'
    Machine{ stack=[] }                 -> hcf m' "IF: stack underflow"
    Machine{ stack=StackBool False:xs } -> return m'{ stack=xs }
    Machine{ stack=StackBool True:xs }  -> evalBlocks [b] m'{ stack=xs }
    Machine{ stack=_:_ }                -> hcf m' "IF: type mis-match"

evalBlocks :: [Block] -> Machine -> IO Machine
evalBlocks _ m@Machine{ ok=False } = return m
evalBlocks [] m                    = return m
evalBlocks (BLinear b:bs) m        = applyLinear m b >>= evalBlocks bs
evalBlocks (BIf c b:bs) m          = applyIf c b m >>= evalBlocks bs

interpret :: Program -> IO ()
interpret p = evalBlocks p newMachine >> return ()

interpretFromString :: String -> IO ()
interpretFromString = maybe err interpret . stringToProgram
  where
    err = putStrLn "Unable to parse program"

interpretFromFile :: FilePath -> IO ()
interpretFromFile path = readFile path >>= interpretFromString