module AssemblerInterpreter where

import Debug.Trace

import Data.Function((&))
import Data.Maybe (fromMaybe)
import Text.Parsec hiding (State)
import Text.Parsec.Char
import qualified Data.Map as M

import Control.Monad.Trans.Maybe
import Control.Monad.Trans.State
import Control.Monad.Trans.Class(lift)

interpret :: String -> Maybe String
interpret prog = Nothing

type Parser = Parsec String ()

data Op =
  Mov | Inc | Dec | Add | Sub  | Mul | Div | Jmp | Cmp | Jne | Je |
  Jge | Jg  | Jle | Jl  | Call | Ret | Msg | End 
  deriving (Show)

opStrP :: Parser String
opStrP = choice $ map string' [
  "mov", "inc", "dec", "add", "sub", "mul", "div", "jmp", "cmp", "jne",
  "je", "jge", "jg", "jle", "jl", "call", "ret", "msg", "end"]

opP :: Parser Op
opP = do
  s <- opStrP
  case s of
    "mov" -> return Mov
    "inc" -> return Inc
    "dec" -> return Dec
    "add" -> return Add
    "sub" -> return Sub
    "mul" -> return Mul
    "div" -> return Div
    "jmp" -> return Jmp
    "cmp" -> return Cmp
    "jne" -> return Jne
    "je"  -> return Je
    "jge" -> return Jge
    "jg"  -> return Jg
    "jle" -> return Jle
    "jl"  -> return Jl
    "call"-> return Call
    "ret" -> return Ret
    "msg" -> return Msg
    "end" -> return End

newtype Label = Label String
  deriving (Show)

inlineSpace :: Parser Char
inlineSpace = oneOf [' ', '\t'] 
inlineSpaces = skipMany $ oneOf [' ', '\t'] 

identifierP :: Parser String
identifierP = do
  notFollowedBy opStrP
  x <- letter <|> oneOf ['_']
  xs <- many (alphaNum <|> oneOf ['_'])
  return (x:xs)

labelP :: Parser Label
labelP = do
  id <- identifierP
  oneOf [':']
  return $ Label id

data Stmt = LabelStmt Label | InstrStmt Op [Arg]
  deriving (Show)

data Arg = IdentifierArg String | NumberArg Int | StringArg String
  deriving (Show)

argP :: Parser Arg
argP = do
  inlineSpaces
  let identifierArgP = IdentifierArg <$> identifierP
      numberP = NumberArg . read <$> many1 digit
      stringP = do
        oneOf ['\''] :: Parser Char
        str <- many $ noneOf ['\'']
        oneOf ['\'']
        return $ StringArg str
  arg <- identifierArgP <|> numberP <|> stringP
  inlineSpaces
  return arg

stmtP :: Parser Stmt
stmtP = do
  let
    labelStmtP = LabelStmt <$> labelP
    argsP = do
      inlineSpaces
      oneOf [',']
      inlineSpaces
      arg <- argP
      inlineSpaces
      return arg
    instrStmtP = do
      op <- opP
      inlineSpaces
      marg <- optionMaybe argP
      args <- case marg of
        Just arg -> do
          targs <- many argsP
          return (arg:targs)
        Nothing -> return []
      return $ InstrStmt op args
  inlineSpaces
  stmt <- labelStmtP <|> instrStmtP
  inlineSpaces
  return stmt

commentP :: Parser ()
commentP = do
  oneOf [';']
  many $ noneOf ['\n']
  return ()

stmtLineP :: Parser [Stmt]
stmtLineP = do
  inlineSpaces
  ms <- optionMaybe stmtP
  inlineSpaces
  optional commentP
  inlineSpaces
  case ms of
    Just s -> return [s]
    Nothing -> return []

progP :: Parser [Stmt]
progP = do
  let newlineP = do
        newline :: Parser Char
        return (++)
  stmts <- chainl stmtLineP newlineP []
  spaces
  eof
  return stmts

processLabel :: [Stmt] -> ([Stmt], M.Map String Int)
processLabel stmts = go stmts [] M.empty 0 where
  go [] processed labelMap i = (reverse processed, labelMap)
  go (instr@(InstrStmt op args):xs) processed labelMap i =
    go xs (instr:processed) labelMap (i+1)
  go ((LabelStmt (Label ident)):xs) processed labelMap i =
    let newLabelMap = M.insert ident i labelMap 
    in
      go xs processed newLabelMap i

data MachineState = MachineState
  { machineRegisters :: M.Map String Int
  , machineStack :: [Int]
  , machineProgCnt :: Int
  , machineOutput :: Maybe String
  , machineProg :: [Stmt]
  , machineLabels :: M.Map String Int
  , machineEnd :: Bool
  }
  deriving (Show)

type MachineM = MaybeT (State MachineState)

nextInstr :: MachineM ()
nextInstr = do
  ms <- lift get
  lift $ put (ms {machineProgCnt = machineProgCnt ms + 1})

fetchInstr :: MachineM (Maybe Stmt)
fetchInstr = do
  ms <- lift get
  let pc = machineProgCnt ms
      prog = machineProg ms
  if pc >= length prog || pc < 0 || machineEnd ms then return Nothing
  else return $ Just (prog !! pc)

setReg :: String -> Int -> MachineM ()
setReg r n = do
  ms <- lift get
  let regs = machineRegisters ms
  lift $ put ms {machineRegisters = M.insert r n regs}

getReg :: String -> MachineM Int
getReg r = do
  ms <- lift get
  let regs = machineRegisters ms
  case M.lookup r regs of
    Just n -> return n
    Nothing -> fail []


execInstr :: Stmt -> MachineM ()
execInstr (InstrStmt op args) = decodeOp op args
execInstr _ = fail []

getResult :: MachineM String
getResult = do
  ms <- lift get
  if not (machineEnd ms) then fail []
  else case machineOutput ms of
    Nothing -> fail []
    Just msg -> return msg

runMachine :: MachineM String
runMachine = do
  ms <- lift get
  instr <- fetchInstr
  case instr of
    Nothing -> getResult
    Just stmt -> do execInstr stmt ; runMachine

execMov :: [Arg] -> MachineM ()
execMov [IdentifierArg reg, NumberArg n] = do setReg reg n
execMov [IdentifierArg r1, IdentifierArg r2] = do
  n <- getReg r1
  setReg r2 n
  nextInstr
execMov _ = fail []

execMsg :: [Arg] -> MachineM ()
execMsg [] = fail []
execMsg args = go args "" where
  go [] msg = do
    ms <- lift get
    lift $ put ms { machineOutput = Just msg }
    nextInstr
  go ((NumberArg n):xs) msg = go  xs (msg ++ show n)
  go ((StringArg s):xs) msg = go xs (msg ++ s)
  go ((IdentifierArg reg):xs) msg = do
    n <- getReg reg
    go xs (msg ++ show n)

execEnd :: [Arg] -> MachineM ()
execEnd [] = do
  ms <- lift get
  lift $ put ms { machineEnd = True }
execEdn _ = fail []

decodeOp :: Op -> [Arg] -> MachineM ()
decodeOp Mov = execMov
decodeOp Msg = execMsg
decodeOp End = execEnd

buildMachine :: String -> Maybe MachineState
buildMachine code = 
  let
    parseResult = runParser progP () "" code
  in
    case parseResult of
      Right lines ->
        let (stmts, labels) = processLabel lines
        in
          Just MachineState
            { machineRegisters = M.empty
            , machineEnd = False
            , machineLabels = labels
            , machineOutput = Nothing
            , machineProg = stmts
            , machineProgCnt = 0
            , machineStack = []
            }
      _ -> Nothing


interp :: String -> Maybe String
interp code = do
  ms <- buildMachine code
  let (res, _) = runState (runMaybeT runMachine) ms
  res

main = do
  print $ runParser opP () "" "msg" 
  print $ runParser stmtLineP () "" "ret" 
  print $ runParser progP () "" "ret \n  mov a , 12\n uwu: \n ret" 
  print $ runParser progP () "" "ret ; 123 ; 123\n  mov a , 12\n uwu: \n ret" 
  print $ interp "msg \'hello\'\n end"