module Lib
    ( tokenize, BrainFuckToken, parseTokens, BrainFuckOperation, compile, CompileState
    ) where

import Control.Monad (liftM)
import qualified Data.Word as Word
import Data.Bits
import Data.Int (Int32)

data BrainFuckToken = OperatorPlus Int | OperatorMinus Int | OperatorLeft Int | OperatorRight Int | OperatorSkip | OperatorContinue | OperatorPut | OperatorRead | Comment String
    deriving Show

partitionOn :: (a -> Bool) -> [a] -> ([a], [a])
partitionOn condition (a:as)
    | condition a = ([], a:as)
    | otherwise   = (a:pre, post)
    where
        (pre, post) = partitionOn condition as
partitionOn _ []  = ([], [])

tokenizeQuantified :: Char -> String -> (Int -> BrainFuckToken) -> [BrainFuckToken]
tokenizeQuantified target cs constructor = constructor (length targetPartition) : tokenize restPartition
    where
        (targetPartition, restPartition) = partitionOn (/=target) cs

tokenizeComment :: String -> [BrainFuckToken]
tokenizeComment s = Comment comment : tokenize rest
    where
        (comment, rest) = partitionOn (flip elem "+-<>[]") s

tokenize :: String -> [BrainFuckToken]
tokenize []       = []
tokenize (c:cs)
    | c == '+'  = tokenizeQuantified '+' (c:cs) OperatorPlus
    | c == '-'  = tokenizeQuantified '-' (c:cs) OperatorMinus
    | c == '<'  = tokenizeQuantified '<' (c:cs) OperatorLeft
    | c == '>'  = tokenizeQuantified '>' (c:cs) OperatorRight
    | c == '['  = OperatorSkip     : tokenize cs
    | c == ']'  = OperatorContinue : tokenize cs
    | c == '.'  = OperatorPut      : tokenize cs
    | c == ','  = OperatorRead     : tokenize cs
    | otherwise = tokenizeComment (c:cs)

data BrainFuckOperation = IncrementCell Int | DecrementCell Int | MoveLeft Int | MoveRight Int | PutChar | ReadChar | Conditionally [BrainFuckOperation]
    deriving Show

parseTokens :: [BrainFuckToken] -> Either String [BrainFuckOperation]
parseTokens ts = liftM (fst) (parseTokens' 0 ts)

parseTokens' :: Int -> [BrainFuckToken] -> Either String ([BrainFuckOperation], [BrainFuckToken])
parseTokens' 0 []                        = Right ([], [])
parseTokens' n []                        = Left ("Missing " ++ show n ++ " closing brackets ']'")
parseTokens' l (OperatorPlus  amount:ts) = liftM (\t -> (IncrementCell amount:fst t, snd t)) (parseTokens' l ts)
parseTokens' l (OperatorMinus amount:ts) = liftM (\t -> (DecrementCell amount:fst t, snd t)) (parseTokens' l ts)
parseTokens' l (OperatorLeft  amount:ts) = liftM (\t -> (MoveLeft      amount:fst t, snd t)) (parseTokens' l ts)
parseTokens' l (OperatorRight amount:ts) = liftM (\t -> (MoveRight     amount:fst t, snd t)) (parseTokens' l ts)
parseTokens' l (OperatorSkip        :ts) = either (Left) (\(cs, _) -> either (Left) (Right . \(os, rs) -> (Conditionally cs:os, rs)) r) t
    where
        t = parseTokens' (l+1) ts
        r = either (Left) (parseTokens' l . snd) t
parseTokens' 0 (OperatorContinue    :_)  = Left ("Found ] when there was none expected")
parseTokens' _ (OperatorContinue    :ts) = Right ([], ts)
parseTokens' l (OperatorPut         :ts) = liftM (\t -> ( PutChar:fst t, snd t)) (parseTokens' l ts)
parseTokens' l (OperatorRead        :ts) = liftM (\t -> (ReadChar:fst t, snd t)) (parseTokens' l ts)
parseTokens' l (Comment _           :ts) = parseTokens' l ts


ubcOpcodeLoadUpper :: Word.Word32
ubcOpcodeLoadUpper = 0x02000000

ubcOpcodeLoadLower :: Word.Word32
ubcOpcodeLoadLower = 0x03000000

ubcOpcodeMem2Reg :: Word.Word32
ubcOpcodeMem2Reg  = 0x04000000
ubcOpcodeReg2Mem :: Word.Word32
ubcOpcodeReg2Mem  = 0x05000000

ubcOpcodeJumpRel :: Word.Word32
ubcOpcodeJumpRel  = 0x09000000

ubcOpcodeForeignCall :: Word.Word32
ubcOpcodeForeignCall  = 0x0F000000

ubcOpcodeAddU :: Word.Word32
ubcOpcodeAddU  = 0x80000000
ubcOpcodeSubU :: Word.Word32
ubcOpcodeSubU  = 0x81000000
ubcOpcodeCmpU :: Word.Word32
ubcOpcodeCmpU  = 0x89000000

ubcConditionNone :: Word.Word32
ubcConditionNone = 0x00000000
ubcConditionEq   :: Word.Word32
ubcConditionEq   = 0x00800000

ubcForeignCall :: Int32 -> Word.Word32
ubcForeignCall r = ubcOpcodeForeignCall .|. ubcInstructionRegArg1 r

ubcInstructionRegArg1 :: Int32 -> Word.Word32
ubcInstructionRegArg1 r
    | r < 16    = (shiftL (fromIntegral r .&. 0xF) 16)
    | otherwise = error "Register index > 16"

ubcInstructionRegArg2 :: Int32 -> Word.Word32
ubcInstructionRegArg2 r
    | r < 16    = (shiftL (fromIntegral r .&. 0xF) 12)
    | otherwise = error "Register index > 16"

ubcInstructionRegArg3 :: Int32 -> Word.Word32
ubcInstructionRegArg3 r
    | r < 16    = (shiftL (fromIntegral r .&. 0xF) 8)
    | otherwise = error "Register index > 16"

ubcInstructionPayLoad :: Word.Word32 -> Word.Word32
ubcInstructionPayLoad w = 0x0000FFFF .&. w

ubcLoadUpper :: Word.Word16 -> Int32 -> Word.Word32
ubcLoadUpper w r = ubcOpcodeLoadUpper .|. ubcInstructionRegArg1 r .|. ubcInstructionPayLoad (fromIntegral w)

ubcLoadLower :: Word.Word16 -> Int32 -> Word.Word32
ubcLoadLower w r = ubcOpcodeLoadLower .|. ubcInstructionRegArg1 r .|. ubcInstructionPayLoad (fromIntegral w)

ubcLoadIntWords :: Int32 -> Int32 -> [Word.Word32]
ubcLoadIntWords i r = [ubcLoadUpper (fromIntegral u) r, ubcLoadLower (fromIntegral l) r]
    where
        w = fromIntegral i :: Word.Word32
        u = w .>>. 16
        l = w .&. 0xFFFF

ubcCmpU :: Int32 -> Int32 -> Word.Word32
ubcCmpU a b = ubcOpcodeCmpU .|. ubcInstructionRegArg1 a .|. ubcInstructionRegArg2 b

ubcJumpRel :: Int32 -> Int32 -> Word.Word32 -> [Word.Word32]
ubcJumpRel d r c = ubcLoadIntWords d r ++ [ubcOpcodeJumpRel .|. ubcInstructionRegArg1 r .|. c]

ubcLoadCellWord :: Int32 -> Word.Word32
ubcLoadCellWord r = ubcOpcodeMem2Reg .|. ubcConditionNone .|. ubcInstructionRegArg1 r .|. ubcInstructionRegArg2 0

ubcStoreCellWord :: Int32 -> Word.Word32
ubcStoreCellWord r = ubcOpcodeReg2Mem .|. ubcInstructionRegArg1 0 .|. ubcInstructionRegArg2 r

ubcAddWord :: Int32 -> Int32 -> Int32 -> Word.Word32
ubcAddWord d a b = ubcOpcodeAddU .|. ubcInstructionRegArg1 d .|. ubcInstructionRegArg2 a .|. ubcInstructionRegArg3 b

ubcSubWord :: Int32 -> Int32 -> Int32 -> Word.Word32
ubcSubWord d a b = ubcOpcodeSubU .|. ubcInstructionRegArg1 d .|. ubcInstructionRegArg2 a .|. ubcInstructionRegArg3 b

ubcIncrementRegisterWords :: Int32 -> Int32 -> Int32 -> [Word.Word32]
ubcIncrementRegisterWords i ir cr = ubcLoadIntWords i ir ++ [ubcAddWord cr cr ir]

ubcDecrementRegisterWords :: Int32 -> Int32 -> Int32 -> [Word.Word32]
ubcDecrementRegisterWords i ir cr = ubcLoadIntWords i ir ++ [ubcSubWord cr cr ir]

data CompileState = CompileState { isCellLoaded :: Bool, isCellModified :: Bool }

compileIncrement :: CompileState -> Int32 -> [BrainFuckOperation] -> [Word.Word32]
compileIncrement state i os
    | isCellLoaded state = ubcIncrementRegisterWords i 2 1 ++ compile' (state{isCellModified=True}) os
    | otherwise          = ubcLoadCellWord 1 : compileIncrement (state{isCellLoaded=True}) i os

compileDecrement :: CompileState -> Int32 -> [BrainFuckOperation] -> [Word.Word32]
compileDecrement state i os
    | isCellLoaded state = ubcDecrementRegisterWords i 2 1 ++ compile' (state{isCellModified=True}) os
    | otherwise          = ubcLoadCellWord 1 : compileIncrement (state{isCellLoaded=True}) i os

compileMoveLeft :: CompileState -> Int32 -> [BrainFuckOperation] -> [Word.Word32]
compileMoveLeft state d os
    | isCellModified state = ubcStoreCellWord 1 : compileMoveLeft (state{isCellModified=False}) d os
    | otherwise            = ubcDecrementRegisterWords d 2 0 ++ compile' (state{isCellLoaded=False}) os

compileMoveRight :: CompileState -> Int32 -> [BrainFuckOperation] -> [Word.Word32]
compileMoveRight state d os
    | isCellModified state = ubcStoreCellWord 1 : compileMoveLeft (state{isCellModified=False}) d os
    | otherwise            = ubcIncrementRegisterWords d 2 0 ++ compile' (state{isCellLoaded=False}) os

compileConditionally :: CompileState -> [BrainFuckOperation] -> [BrainFuckOperation] -> [Word.Word32]
compileConditionally s cs os
    | isCellLoaded s = skipBlock ++ contained ++ continueBlock ++ compile os
    | otherwise      = ubcLoadCellWord 1 : compileConditionally (s{isCellLoaded=True}) cs os
    where
        contained        = compile cs
        containedLength  = length contained
        forwardDistance  = containedLength + 4 -- plus one to land after the skip Block
        continueDistance = negate (containedLength + 10)
        -- continueBlock has length 1 + 3 = 4
        continueBlock    = ubcStoreCellWord 1 : ubcJumpRel (fromIntegral continueDistance) 2 ubcConditionNone
        -- skipBlock has length 2 + 1 + 3 = 6
        skipBlock        = ubcLoadIntWords 0 2 ++ [ubcCmpU 1 2] ++ (ubcJumpRel (fromIntegral forwardDistance) 2 ubcConditionEq)

compile :: [BrainFuckOperation] -> [Word.Word32]
compile = compile' CompileState{isCellLoaded=False, isCellModified=False}

-- 0: cell pointer register
-- 1: cell value register
-- 2: temporary register
compile' :: CompileState -> [BrainFuckOperation] -> [Word.Word32]
compile' _     []                        = []
compile' state (IncrementCell amount:os) = compileIncrement     state (fromIntegral amount)   os
compile' state (DecrementCell amount:os) = compileDecrement     state (fromIntegral amount)   os
compile' state (MoveRight   distance:os) = compileMoveRight     state (fromIntegral distance) os
compile' state (MoveLeft    distance:os) = compileMoveLeft      state (fromIntegral distance) os
compile' state (Conditionally    ops:os) = compileConditionally state ops                     os
compile' state (PutChar             :os) = ubcLoadIntWords 0 2 ++ [ubcForeignCall 2] ++ compile' state os
compile' state (ReadChar            :os) = ubcLoadIntWords 1 2 ++ [ubcForeignCall 2] ++ compile' state os