{-# LANGUAGE MultiWayIf        #-}
{-# LANGUAGE OverloadedStrings #-}

module Nand2Tetris.VM
  (
    VMInstruction(..),
    BranchingCmd(..),
    ArithmeticCmd(..),
    AccessType(..),
    Segment(..),
    AnnotatedAsmLine(..),
    parseInstruction,
    parseInstructions,
    compileInstruction,
    compileInstructions,
    defaultCompilerEnv,
    runCompiler,
    runCompilerDef,
    bootstrapPreamble
  ) where

import           Control.Monad.State.Strict (State, evalState, forM, gets,
                                             modify')
import           Data.Functor               (($>))
import           Data.Maybe                 (fromMaybe)
import qualified Data.Text                  as T
import           Nand2Tetris.Error          (Error)
import           Nand2Tetris.Hack           (HackInstruction (A, C),
                                             Imm (Imm, Label),
                                             Jump (JEq, JGt, JLt, JNe, JNone, Jmp),
                                             Reg (RegA, RegD, RegM),
                                             Source (S0, SDAndM, SDMinusA, SDOrM, SDPlusA, SDPlusM, SMMinusD, SNeg, SNeg1, SNot, SReg, SRegMinus1, SRegPlus1),
                                             canonicalRepresentation)
import           Text.Megaparsec            (MonadParsec (try), Parsec, choice,
                                             empty, eof, many, (<|>))
import           Text.Megaparsec.Char       (alphaNumChar, char, letterChar,
                                             space1, string)
import           Text.Megaparsec.Char.Lexer (skipLineComment)
import qualified Text.Megaparsec.Char.Lexer as L

data VMInstruction =
    VMArithmetic ArithmeticCmd
  | VMMemAccess AccessType Segment
  | VMBranching BranchingCmd
  | VMFunction FunctionCmd
  deriving (Show, Eq)

data FunctionCmd =
    FFunction T.Text Int
  | FCall T.Text Int
  | FReturn
  deriving (Show, Eq)

data BranchingCmd =
    BLabel T.Text
  | BGoto T.Text
  | BIfGoto T.Text
  deriving (Show, Eq)

data ArithmeticCmd =
    ACAdd
  | ACSub
  | ACNeg
  | ACEq
  | ACGt
  | ACLt
  | ACAnd
  | ACOr
  | ACNot
  deriving (Show, Eq)

data AccessType =
    ACPush
  | ACPop
  deriving (Show, Eq)

--newtype SegmentOffset = SegmentOffset { unSegmentOffset :: Int }
--  deriving (Show, Eq)

data Segment =
    SLocal Int
  | SArgument Int
  | SThis Int
  | SThat Int
  | SConstant Int
  | SStatic T.Text
  | SPointer Int
  | STemp Int
  deriving (Show, Eq)

instrAnnotation :: VMInstruction -> T.Text
instrAnnotation (VMArithmetic ACAdd) = "add"
instrAnnotation (VMArithmetic ACSub) = "sub"
instrAnnotation (VMArithmetic ACNeg) = "neg"
instrAnnotation (VMArithmetic ACEq) = "eq"
instrAnnotation (VMArithmetic ACGt) = "gt"
instrAnnotation (VMArithmetic ACLt) = "lt"
instrAnnotation (VMArithmetic ACAnd) = "and"
instrAnnotation (VMArithmetic ACOr) = "or"
instrAnnotation (VMArithmetic ACNot) = "not"
instrAnnotation (VMMemAccess accessType segment) =
  strAccessType <> " " <> strSegment
  where
    strAccessType = case accessType of
      ACPush -> "push"
      ACPop  -> "pop"

    strSegment = case segment of
      SLocal i      -> "local " <> (T.pack . show $ i)
      SArgument i   -> "argument " <> (T.pack . show $ i)
      SThis i       -> "this " <> (T.pack . show $ i)
      SThat i       -> "that " <> (T.pack . show $ i)
      SConstant i   -> "constant " <> (T.pack . show $ i)
      SStatic label -> "static " <> label
      SPointer i    -> "pointer " <> (T.pack . show $ i)
      STemp i       -> "temp " <> (T.pack . show $ i)
instrAnnotation (VMBranching (BLabel label)) = "label " <> label
instrAnnotation (VMBranching (BGoto label)) = "goto " <> label
instrAnnotation (VMBranching (BIfGoto label)) = "if-goto " <> label
instrAnnotation (VMFunction (FFunction name nvars)) = "function " <> name <> " " <> (T.pack . show $ nvars)
instrAnnotation (VMFunction (FCall name nargs)) = "call " <> name <> " " <> (T.pack . show $ nargs)
instrAnnotation (VMFunction FReturn) = "return"

type Parser = Parsec Error T.Text

parseInstructions :: Parser [VMInstruction]
parseInstructions = many (sp *> parseInstruction <* sp) <* eof
  where
    sp :: Parser ()
    sp = L.space space1 (skipLineComment "//") empty

parseInstruction :: Parser VMInstruction
parseInstruction =
  try parseMemAccess <|>
  try parseArithmetic <|>
  try parseBranching <|>
      parseFunction

  where
    parseFunction :: Parser VMInstruction
    parseFunction = VMFunction <$>
      choice [ string "function" *> space1 *> (FFunction <$> (parseId <* space1) <*> L.decimal),
               string "call" *> space1 *> (FCall <$> (parseId <* space1) <*> L.decimal),
               string "return" $> FReturn
             ]

    parseBranching :: Parser VMInstruction
    parseBranching = VMBranching <$>
      choice [ string "label" *> space1 *> (BLabel <$> parseId),
               string "goto" *> space1 *> (BGoto <$> parseId),
               string "if-goto" *> space1 *> (BIfGoto <$> parseId)
             ]

    parseMemAccess :: Parser VMInstruction
    parseMemAccess = VMMemAccess <$>
      parseAccessType <*>
      parseSegment

    parseAccessType :: Parser AccessType
    parseAccessType = try (string "push" >> space1 $> ACPush) <|>
      (string "pop" >> space1 $> ACPop)

    parseSegment :: Parser Segment
    parseSegment =
      choice [ parseSegmentInt "local" SLocal,
               parseSegmentInt "argument"  SArgument,
               parseSegmentInt "this"  SThis,
               parseSegmentInt "that"  SThat,
               parseSegmentInt "constant" SConstant,
               (string "static" *> space1) *> (SStatic <$> (try parseId <|> parseStaticOffset)),
               parseSegmentInt "pointer" SPointer,
               parseSegmentInt "temp" STemp
             ]


    parseId :: Parser T.Text
    parseId = do
      h <- letterChar
      hs <- many (alphaNumChar <|> char '.' <|> char '_')
      return $ T.pack (h : hs)

    parseStaticOffset :: Parser T.Text
    parseStaticOffset = do
      h <- L.decimal
      return $ "STATIC_" <> (T.pack . show) h

    parseSegmentInt :: T.Text -> (Int -> Segment) -> Parser Segment
    parseSegmentInt sname f = try $ do
      string sname
      space1
      f <$> L.decimal

    parseArithmetic :: Parser VMInstruction
    parseArithmetic = VMArithmetic <$>
      choice [ string "add" $> ACAdd,
               string "sub" $> ACSub,
               string "neg" $> ACNeg,
               string "eq" $> ACEq,
               string "gt" $> ACGt,
               string "lt" $> ACLt,
               string "and" $> ACAnd,
               string "or" $> ACOr,
               string "not" $> ACNot
             ]

data AnnotatedAsmLine =
    Comment T.Text
  | Code HackInstruction
  | ALabel T.Text
  deriving (Show, Eq)

data CompilerEnv = CompilerEnv
  {
    ceLabelCounter    :: Int,
    ceCurrentFunction :: Maybe T.Text
  }

type Compiler = State CompilerEnv

nextLabelId :: Compiler T.Text
nextLabelId = do
  v <- gets ceLabelCounter
  modify' (\s -> s { ceLabelCounter = v + 1 })
  return $ "_VM_LABEL_" <> (T.pack . show) v


defaultCompilerEnv :: CompilerEnv
defaultCompilerEnv = CompilerEnv 0 Nothing

runCompiler :: CompilerEnv -> Compiler a -> a
runCompiler = flip evalState

runCompilerDef :: Compiler a -> a
runCompilerDef = runCompiler defaultCompilerEnv

compileInstructions :: [VMInstruction] -> [AnnotatedAsmLine]
compileInstructions instrs = concat $ runCompilerDef (forM instrs compileInstruction)

compileInstruction :: VMInstruction -> Compiler [AnnotatedAsmLine]
compileInstruction instr = do
  let ann = Comment (instrAnnotation instr)
  compiled <- compileInstruction' instr
  return $ ann : compiled

  where
    compileInstruction' :: VMInstruction -> Compiler [AnnotatedAsmLine]
    compileInstruction' (VMMemAccess ACPush segment) = return $
      case segment of
        SLocal offset -> pushSeq "LCL" offset
        SArgument offset -> pushSeq "ARG" offset
        SThis offset -> pushSeq "THIS" offset
        SThat offset -> pushSeq "THAT" offset
        SConstant offset -> [
          Code (A . Imm $ offset),
          Code (C [RegD] (SReg RegA) JNone),
          Code (A (Label "SP")),
          Code (C [RegA] (SReg RegM) JNone),
          Code (C [RegM] (SReg RegD) JNone),
          Code (C [RegD] (SRegPlus1 RegA) JNone),
          Code (A (Label "SP")),
          Code (C [RegM] (SReg RegD) JNone) ]
        SStatic label -> pushLabelValue label
        SPointer offset -> if offset == 0 then pushLabelValue "THIS" else pushLabelValue "THAT"
        STemp offset -> [
          Code (A . Imm $ (5 + offset)),
          Code (C [RegD] (SReg RegM) JNone),
          Code (A (Label "SP")),
          Code (C [RegA] (SReg RegM) JNone),
          Code (C [RegM] (SReg RegD) JNone),
          Code (C [RegD] (SRegPlus1 RegA) JNone),
          Code (A (Label "SP")),
          Code (C [RegM] (SReg RegD) JNone) ]

    compileInstruction' (VMMemAccess ACPop segment) = return $
      case segment of
        SLocal offset -> popSeq "LCL" offset
        SArgument offset -> popSeq "ARG" offset
        SThis offset -> popSeq "THIS" offset
        SThat offset -> popSeq "THAT" offset
        SConstant _ -> error "Pop constant instruction"
        SStatic label -> popLabelValue label
        SPointer offset -> if offset == 0 then popLabelValue "THIS" else popLabelValue "THAT"
        STemp offset -> popRam (5 + offset)

    compileInstruction' (VMArithmetic ACAdd) = return $ binopSequence SDPlusM
    compileInstruction' (VMArithmetic ACSub) = return $ binopSequence SMMinusD
    compileInstruction' (VMArithmetic ACNeg) = return $ unopSequence (SNeg RegM)
    compileInstruction' (VMArithmetic ACEq) = compileConditional JEq
    compileInstruction' (VMArithmetic ACGt) = compileConditional JGt
    compileInstruction' (VMArithmetic ACLt) = compileConditional JLt
    compileInstruction' (VMArithmetic ACAnd) = return $ binopSequence SDAndM
    compileInstruction' (VMArithmetic ACOr) = return $ binopSequence SDOrM
    compileInstruction' (VMArithmetic ACNot) = return $ unopSequence (SNot RegM)
    compileInstruction' (VMBranching (BLabel label)) = do
      fname <- fromMaybe "_Global" <$> gets ceCurrentFunction
      return [ALabel $ fname <> "$" <> label]
    compileInstruction' (VMBranching (BGoto label)) = do
      fname <- fromMaybe "_Global" <$> gets ceCurrentFunction
      return [ Code (A (Label $ fname <> "$" <> label)),
               Code (C [] S0 Jmp)
             ]
    compileInstruction' (VMBranching (BIfGoto label)) = do
      fname <- fromMaybe "_Global" <$> gets ceCurrentFunction
      return [ Code (A (Label "SP")),
               Code (C [RegA, RegM] (SRegMinus1 RegM) JNone),
               Code (C [RegD] (SReg RegM) JNone),
               Code (A (Label $ fname <> "$" <> label)),
               Code (C [] (SReg RegD) JNe)
             ]
    compileInstruction' (VMFunction (FFunction name nvars)) = do
      modify' (\s -> s { ceCurrentFunction = Just name })
      return ([ ALabel name,
                Code (A (Label "SP")),
                Code (C [RegA] (SReg RegM) JNone)
                ] <>
               (concat . replicate nvars $
                 [ Code (C [RegM] S0 JNone),
                   Code (C [RegA] (SRegPlus1 RegA) JNone)
                 ]) <>
             [ Code (C [RegD] (SReg RegA) JNone),
               Code (A (Label "SP")),
               Code (C [RegM] (SReg RegD) JNone)
             ]
             )

    compileInstruction' (VMFunction (FCall name nargs)) = do
      lab <- nextLabelId
      mbCurrentFunction <- gets ceCurrentFunction
      let labelId =
            case mbCurrentFunction of
              Just fname -> fname <> ".ret." <> lab
              Nothing    -> lab
      return $ pushLabel labelId <>
               pushLabelValue "LCL" <>
               pushLabelValue "ARG" <>
               pushLabelValue "THIS" <>
               pushLabelValue "THAT" <>
               -- ARG = SP - 5 - nargs
               [ Code (A (Label "SP")),
                 Code (C [RegD] (SReg RegM) JNone),
                 Code (A (Imm 5)),
                 Code (C [RegD] SDMinusA JNone),
                 Code (A (Imm nargs)),
                 Code (C [RegD] SDMinusA JNone),
                 Code (A (Label "ARG")),
                 Code (C [RegM] (SReg RegD) JNone)
               ] <>
               -- LCL = SP
               [ Code (A (Label "SP")),
                 Code (C [RegD] (SReg RegM) JNone),
                 Code (A (Label "LCL")),
                 Code (C [RegM] (SReg RegD) JNone)
               ] <>
               -- goto function & and generate return label
               [
                 Code (A (Label name)),
                 Code (C [] S0 Jmp),
                 ALabel labelId
               ]

    compileInstruction' (VMFunction FReturn) = do
      return $
        -- endFrame (R15) = LCL
        [ Code (A (Label "LCL")),
          Code (C [RegD] (SReg RegM) JNone),
          Code (A (Imm 15)),
          Code (C [RegM] (SReg RegD) JNone)
        ] <>
        -- retAddr (R14) = *(endFrame - 5)
        [ Code (A (Imm 5)),
          Code (C [RegA] SDMinusA JNone),
          Code (C [RegD] (SReg RegM) JNone),
          Code (A (Imm 14)),
          Code (C [RegM] (SReg RegD) JNone)
        ] <>
        --  *ARG = pop()
        [ Code (A (Label "SP")),
          Code (C [RegA] (SRegMinus1 RegM) JNone),
          Code (C [RegD] (SReg RegM) JNone),
          Code (A (Label "ARG")),
          Code (C [RegA] (SReg RegM) JNone),
          Code (C [RegM] (SReg RegD) JNone)
        ] <>
        --  SP = ARG+1
        [ Code (C [RegD] (SRegPlus1 RegA) JNone),
          Code (A (Label "SP")),
          Code (C [RegM] (SReg RegD) JNone)
        ] <>
        -- THAT = *(endFrame - 1), endFrame--
        [ Code (A (Imm 15)),
          Code (C [RegA, RegM] (SRegMinus1 RegM) JNone),
          Code (C [RegD] (SReg RegM) JNone),
          Code (A (Label "THAT")),
          Code (C [RegM] (SReg RegD) JNone)
        ] <>
        -- THIS = *(endFrame - 1), endFrame--
        [ Code (A (Imm 15)),
          Code (C [RegA, RegM] (SRegMinus1 RegM) JNone),
          Code (C [RegD] (SReg RegM) JNone),
          Code (A (Label "THIS")),
          Code (C [RegM] (SReg RegD) JNone)
        ] <>
        -- ARG = *(endFrame - 1), endFrame--
        [ Code (A (Imm 15)),
          Code (C [RegA, RegM] (SRegMinus1 RegM) JNone),
          Code (C [RegD] (SReg RegM) JNone),
          Code (A (Label "ARG")),
          Code (C [RegM] (SReg RegD) JNone)
        ] <>
        -- ARG = *(endFrame - 1), endFrame--
        [ Code (A (Imm 15)),
          Code (C [RegA, RegM] (SRegMinus1 RegM) JNone),
          Code (C [RegD] (SReg RegM) JNone),
          Code (A (Label "LCL")),
          Code (C [RegM] (SReg RegD) JNone)
        ] <>
        -- goto retAddr
        [ Code (A (Imm 14)),
          Code (C [RegA] (SReg RegM) JNone),
          Code (C [] S0 Jmp)
        ]

    compileConditional jumpType = do
      let sub = binopSequence SMMinusD
      labEq <- nextLabelId
      labDone <- nextLabelId
      let rest = [
            Code (A (Label "SP")),
            Code (C [RegA] (SRegMinus1 RegM) JNone),
            Code (C [RegD] (SReg RegM) JNone),
            Code (A (Label labEq)),
            Code (C [] (SReg RegD) jumpType),
            Code (A (Label "SP")),
            Code (C [RegA] (SRegMinus1 RegM) JNone),
            Code (C [RegM] S0 JNone),
            Code (A (Label labDone)),
            Code (C [] S0 Jmp),

            ALabel labEq,

            Code (A (Label "SP")),
            Code (C [RegA] (SRegMinus1 RegM) JNone),
            Code (C [RegM] SNeg1 JNone),

            ALabel labDone
            ]
      return $ sub ++ rest

    unopSequence cmd =
      [
        Code (A (Label "SP")),
        Code (C [RegA] (SRegMinus1 RegM) JNone),
        Code (C [RegM] cmd JNone)
      ]

    binopSequence cmd =
      [
        Code (A (Label "SP")),
        Code (C [RegA] (SRegMinus1 RegM) JNone),
        Code (C [RegD] (SReg RegM) JNone),
        Code (C [RegA] (SRegMinus1 RegA) JNone),
        Code (C [RegM] cmd JNone),
        Code (C [RegD] (SRegPlus1 RegA) JNone),
        Code (A (Label "SP")),
        Code (C [RegM] (SReg RegD) JNone)
      ]

    pushLabel label = [
          Code (A . Label $ label),
          Code (C [RegD] (SReg RegA) JNone),
          Code (A (Label "SP")),
          Code (C [RegA] (SReg RegM) JNone),
          Code (C [RegM] (SReg RegD) JNone),
          Code (C [RegD] (SRegPlus1 RegA) JNone),
          Code (A (Label "SP")),
          Code (C [RegM] (SReg RegD) JNone) ]

    pushLabelValue label = [
          Code (A . Label $ label),
          Code (C [RegD] (SReg RegM) JNone),
          Code (A (Label "SP")),
          Code (C [RegA] (SReg RegM) JNone),
          Code (C [RegM] (SReg RegD) JNone),
          Code (C [RegD] (SRegPlus1 RegA) JNone),
          Code (A (Label "SP")),
          Code (C [RegM] (SReg RegD) JNone) ]

    pushSeq regName offset = [
          Code (A . Imm $ offset),
          Code (C [RegD] (SReg RegA) JNone),
          Code (A (Label regName)),
          Code (C [RegA] (SReg RegM) JNone),
          Code (C [RegA] SDPlusA JNone),
          Code (C [RegD] (SReg RegM) JNone),
          Code (A (Label "SP")),
          Code (C [RegA] (SReg RegM) JNone),
          Code (C [RegM] (SReg RegD) JNone),
          Code (C [RegD] (SRegPlus1 RegA) JNone),
          Code (A (Label "SP")),
          Code (C [RegM] (SReg RegD) JNone) ]


    popSeq regName offset = [
                           Code (A (Label "SP")),
                           Code (C [RegM] (SRegMinus1 RegM) JNone),

                           Code (A (Imm offset)),
                           Code (C [RegD] (SReg RegA) JNone),
                           Code (A (Label regName)),
                           Code (C [RegD] SDPlusM JNone),
                           Code (A (Imm 15)),
                           Code (C [RegM] (SReg RegD) JNone),

                           Code (A (Label "SP")),
                           Code (C [RegA] (SReg RegM) JNone),
                           Code (C [RegD] (SReg RegM) JNone),
                           Code (A (Imm 15)),
                           Code (C [RegA] (SReg RegM) JNone),
                           Code (C [RegM] (SReg RegD) JNone)
                            ]
    popLabelValue label = [
                           Code (A (Label "SP")),
                           Code (C [RegM] (SRegMinus1 RegM) JNone),
                           Code (C [RegA] (SReg RegM) JNone),
                           Code (C [RegD] (SReg RegM) JNone),

                           Code (A . Label $ label),
                           Code (C [RegM] (SReg RegD) JNone) ]
    popRam offset = [
                           Code (A (Label "SP")),
                           Code (C [RegM] (SRegMinus1 RegM) JNone),
                           Code (C [RegA] (SReg RegM) JNone),
                           Code (C [RegD] (SReg RegM) JNone),

                           Code (A . Imm $ offset),
                           Code (C [RegM] (SReg RegD) JNone) ]

bootstrapPreamble :: [AnnotatedAsmLine]
bootstrapPreamble = [ Comment "Bootstrap preamble",
                      Code (A (Imm 261)),
                      Code (C [RegD] (SReg RegA) JNone),
                      Code (A (Label "SP")),
                      Code (C [RegM] (SReg RegD) JNone),
                      Code (A (Label "LCL")),
                      Code (C [RegM] (SReg RegD) JNone),
                      Code (A (Label "Sys.init")),
                      Code (C [] S0 Jmp)
                    ]