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

module Nand2Tetris.VM
  (
    VMInstruction(..),
    ArithmeticCmd(..),
    AccessType(..),
    Segment(..),
    AnnotatedAsmLine(..),
    parseInstruction,
    compileInstruction
  ) where

import Nand2Tetris.Error ( Error )
import Text.Megaparsec (Parsec, MonadParsec (try), (<|>), choice, many)
import qualified Data.Text as T
import Text.Megaparsec.Char (string, space1, letterChar, alphaNumChar, char)
import Data.Functor (($>))
import qualified Text.Megaparsec.Char.Lexer as L
import Nand2Tetris.Hack (HackInstruction (A, C), Imm (Label, Imm), Reg (RegD, RegA, RegM), Source (SReg, SDPlusA, SRegPlus1, SRegMinus1, SDPlusM), Jump (JNone))

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 i) = ""
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

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 <$> parseId),
               parseSegmentInt "pointer" SPointer,
               parseSegmentInt "temp" STemp
             ]


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


    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
  deriving (Show, Eq)

compileInstruction :: VMInstruction -> [AnnotatedAsmLine]
compileInstruction instr =
  Comment (instrAnnotation instr) : compileInstruction' instr

  where
    compileInstruction' (VMMemAccess ACPush segment) =
      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) =
      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)


    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 [RegA] SDPlusM JNone),
                           Code (C [RegD] (SReg RegM) JNone),

                           Code (A (Label "SP")),
                           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) ]