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

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

import           Control.Monad.State.Strict (State, evalState, forM, gets,
                                             modify')
import           Data.Functor               (($>))
import qualified Data.Text                  as T
import           Nand2Tetris.Error          (Error)
import           Nand2Tetris.Hack           (HackInstruction (A, C),
                                             Imm (Imm, Label),
                                             Jump (JEq, JGt, JLt, JNone, Jmp),
                                             Reg (RegA, RegD, RegM),
                                             Source (S0, SDAndM, 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
  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)


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 <|>
      parseArithmetic

  where
    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
  }

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

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)

    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)
      ]

    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) ]