module Parse
( parse
, generateNASM
) where

import Control.Monad (join)
import qualified Data.Map as Map

data Node = Loop [Node]
          | MoveRight Integer
          | MoveLeft Integer
          | Add Integer
          | Subtract Integer
          | PutChar Integer
          | GetChar Integer
          | Nop Integer -- Used only during parsing/compiling
          | Exit Integer -- XXX: NONSTANDARD
          deriving Show
type Program = Map.Map String [Node]

takeSection :: Eq a => [a] -> (a, a) -> Maybe ([a], [a])
takeSection code (up, down) = rec code [] [] 1
    where rec [] pre post 0 = Just (reverse pre, [])
          rec [] _ _ _ = Nothing
          rec (x:xs) pre post level
              | level == 0 = Just (reverse pre, (x:xs))
              | x == up    = rec xs (x:pre) post (level+1)
              | x == down  = rec xs (x:pre) post (level-1)
              | otherwise  = rec xs (x:pre) post level

parse :: String -> [Node]
parse code = rec code 0 '\0' (Nop)
    where rec [] multiplier _ op
              | multiplier > 0 = (op multiplier) : []
              | otherwise = []
          rec (x:xs) multiplier previous op
              | multiplier > 0 && x /= previous = (op multiplier) : rec (x:xs) 0 previous (Nop)
              | x == previous                   = rec xs (multiplier+1) x op
              | x == '>'                        = rec xs 1 x (MoveRight)
              | x == '<'                        = rec xs 1 x (MoveLeft)
              | x == '+'                        = rec xs 1 x (Add)
              | x == '-'                        = rec xs 1 x (Subtract)
              | x == '.'                        = rec xs 1 x (PutChar)
              | x == ','                        = rec xs 1 x (GetChar)
              | x == '!'                        = rec xs 1 x (Exit)
              | x == '['                        = case loopSections of
                                                       Just (inside, after) -> (Loop $ sub inside) : (sub after)
                                                       Nothing -> error "Unterminated loop"
              | otherwise                       = rec xs multiplier x op
              where loopSections = takeSection xs ('[', ']')
                    notLoopSections = takeSection xs ('(', ')')
                    sub = \x -> rec x 0 '\0' (Nop)

generateNASM :: [String] -> [Node] -> [String]
generateNASM overhead program = overhead ++ rec program [] 0 ++ ["xor rdi, rdi", "call exit"]
    where rec :: [Node] -> [Int] -> Int -> [String]
          rec (MoveRight x:xs) counts count = ("add r9, " ++ (show x)) : (rec xs counts (count+1))
          rec (MoveLeft x:xs) counts count = ("sub r9, " ++ (show x)) : (rec xs counts (count+1))
          rec (Add x:xs) counts count =
              if x == 1
                   then "inc byte [r9]" : (rec xs counts (count+1))
                   else ("add byte [r9], " ++ (show x)) : (rec xs counts (count+1))
          rec (Subtract x:xs) counts count =
              if x == 1
                   then "dec byte [r9]" : (rec xs counts (count+1))
                   else ("sub byte [r9], " ++ (show x)) : (rec xs counts (count+1))
          rec (Nop _:xs) counts count = (rec xs counts (count+1)) -- Skip no-ops
          rec (PutChar x:xs) counts count = ("mov rdx, " ++ (show x)) : "call puts" : (rec xs counts (count+1))
          rec (GetChar x:xs) counts count = ("mov rdx, " ++ (show x)) : "call gets" : (rec xs counts (count+1))
          rec (Exit x:xs) counts count = ["mov rdi, [r9]", "call exit"]
          rec (Loop x:xs) counts count = (startlabel ++ ":")
                                       : "cmp byte [r9], byte 0"
                                       : ("je " ++ endlabel)
                                       : (rec x ((count+1) : counts) (count+1))
                                       ++ ["jmp " ++ startlabel]
                                       ++ [endlabel ++ ":"]
                                       ++ (rec xs (count : counts) (count+1))
                                       where startlabel = ("l" ++ (join $ map show (count : counts)))
                                             endlabel = "e" ++ startlabel
          rec [] counts count = []
