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

module Nand2Tetris.Hack
  (
    HackInstruction(..),
    Imm(..),
    Reg(..),
    Source(..),
    Jump(..),
    AsmLine(..),
    hackParser,
    defaultVars,
    compile,
    LabelMap,
    Error(..),
    addVars,
    addLabels
  ) where

import           Control.Monad              (forM, when)
import           Data.Bits                  (Bits (shift), (.|.))
import           Data.Foldable              (Foldable (foldl'))
import           Data.Functor               (($>))
import qualified Data.Map.Strict            as M
import           Data.Maybe                 (fromMaybe)
import qualified Data.Text                  as T
import           Data.Word                  (Word16)
import           Text.Megaparsec            (MonadParsec (try), Parsec,
                                             ShowErrorComponent, choice, empty,
                                             eof, many, oneOf, optional,
                                             skipManyTill, some, (<|>))
import           Text.Megaparsec.Char       (alphaNumChar, char, letterChar,
                                             space, space1, string, string')
import           Text.Megaparsec.Char.Lexer (decimal, skipLineComment)
import qualified Text.Megaparsec.Char.Lexer as L
import Text.Megaparsec.Error (ShowErrorComponent(showErrorComponent))

defaultVars :: LabelMap
defaultVars = M.fromList $ [("SP", 0), ("LCL", 1), ("ARG", 2), ("THIS", 3), ("THAT", 4)] <>
      zip (fmap (\x -> "R" <> (T.pack . show $ x)) [0..15 :: Int]) [0..15 :: Int]

data HackInstruction =
    A Imm
  | C [Reg] Source Jump
  deriving (Eq, Show)

data Imm = Imm Int | Label T.Text
  deriving (Show, Eq)

data Reg = RegA | RegM | RegD
  deriving (Eq, Show)

data Source =
    S0
  | S1
  | SNeg1
  | SReg Reg
  | SNot Reg
  | SNeg Reg
  | SRegPlus1 Reg
  | SRegMinus1 Reg
  | SDPlusA
  | SDPlusM
  | SDMinusA
  | SDMinusM
  | SAMinusD
  | SMMinusD
  | SDAndA
  | SDAndM
  | SDOrA
  | SDOrM
  deriving (Show, Eq)

data Jump =
    JNone
  | JGt
  | JEq
  | JGe
  | JLt
  | JNe
  | JLe
  | Jmp
  deriving (Show, Eq)

data Error = ParsingError T.Text Int
  deriving (Show, Eq, Ord)

instance ShowErrorComponent Error where
  showErrorComponent = show
type Parser = Parsec Error T.Text

data AsmLine =
    LineInstruction HackInstruction
  | LineLabel T.Text
  deriving (Show, Eq)

hackParser :: Parser [AsmLine]
hackParser = many (sp *> choice [parseInstruction, parseLabel] <* sp) <* eof
  where
    parseInstruction :: Parser AsmLine
    parseInstruction = try (aInstr <|> cInstr)

    aInstr :: Parser AsmLine
    aInstr = do
      char '@'
      space
      aInstrOffset <|> aInstrLabel

    aInstrOffset :: Parser AsmLine
    aInstrOffset = LineInstruction . A . Imm <$> decimal

    aInstrLabel :: Parser AsmLine
    aInstrLabel = LineInstruction . A . Label <$> idParser

    cInstr = do
      target <- optional $ try (many parseReg <* (space >> char '=' >> space))
      comp <- parseComputation
      jump <- optional parseJump

      return . LineInstruction $ C (fromMaybe [] target) comp (fromMaybe JNone jump)

    parseReg :: Parser Reg
    parseReg = try (char 'A' $> RegA) <|> try (char 'D' $> RegD) <|> try (char 'M' $> RegM)

    parseComputation :: Parser Source
    parseComputation = space >>
      (try $ char '0' $> S0) <|>
      (try $ char '1' $> S1) <|>
      (try $ string "-1" $> SNeg1) <|>
      (try $ char '!' *> space *> parseReg >>= pure . SNot) <|>
      (try $ char '-' *> space *> parseReg >>= pure . SNeg) <|>
      (try $ parseReg <* space <* char '+' <* space <* char '1' >>= pure . SRegPlus1) <|>
      (try $ parseReg <* space <* char '-' <* space <* char '1' >>= pure . SRegMinus1) <|>
      (try parseBinop) <|>
      (try $ parseReg >>= pure . SReg)

    parseBinop :: Parser Source
    parseBinop = do
      space
      r1 <- parseReg
      space
      op <- oneOf ['+', '-', '&', '|']
      space
      r2 <- parseReg
      when (r1 == r2) $ fail "Same register is not allowed in binary operation"
      when (r1 /= RegD && r2 /= RegD) $ fail "One of operands of binary operation should be D"
      if
        | op == '+' ->
          return $ if r1 == RegA || r2 == RegA then SDPlusA else SDPlusM
        | op == '&' ->
          return $ if r1 == RegA || r2 == RegA then SDAndA else SDAndM
        | op == '|' ->
          return $ if r1 == RegA || r2 == RegA then SDOrA else SDOrM
        | op == '-' && r1 == RegD -> return $ if r2 == RegA then SDMinusA else SDMinusM
        | op == '-' && r1 == RegA -> return SAMinusD
        | op == '-' && r1 == RegM -> return SMMinusD
        | otherwise -> fail "Invalid operands"

    parseJump :: Parser Jump
    parseJump = try $ do
      space
      char ';'
      space
      (try $ string' "jgt" *> pure JGt) <|>
        (try $ string' "jeq" *> pure JEq) <|>
        (try $ string' "jge" *> pure JGe) <|>
        (try $ string' "jlt" *> pure JLt) <|>
        (try $ string' "jne" *> pure JNe) <|>
        (try $ string' "jle" *> pure JLe) <|>
        (try $ string' "jmp" *> pure Jmp)

    parseLabel :: Parser AsmLine
    parseLabel = try $ do
      char '('
      labelId <- idParser
      char ')'
      return $ LineLabel labelId

    idParser :: Parser T.Text
    idParser = T.pack <$> some (letterChar <|> oneOf ['.', '_', '$']) <> many alphaNumChar

    sp :: Parser ()
    sp = L.space space1 (skipLineComment "//") empty

type LabelMap = M.Map T.Text Int

addLabels :: LabelMap -> [AsmLine] -> LabelMap
addLabels m parsed = snd $ foldl' addLabel (0, m) parsed

addVars :: LabelMap -> LabelMap -> [AsmLine] -> LabelMap
addVars m labels parsed = snd $ foldl' (addVar labels) (16, m) parsed

addLabel :: (Int, LabelMap) -> AsmLine -> (Int, LabelMap)
addLabel (i, m) (LineInstruction _) = (i + 1, m)
addLabel (i, m) (LineLabel v)       = (i, M.insert v i m)

addVar :: LabelMap -> (Int, LabelMap) -> AsmLine -> (Int, LabelMap)
addVar labs (cnt, m) (LineInstruction (A (Label v))) = case M.lookup v labs of
  Just _ -> (cnt, m)
  Nothing -> case M.lookup v m of
    Just _  -> (cnt, m)
    Nothing -> (cnt + 1, M.insert v cnt m)
addVar _ m _ = m

compile :: LabelMap -> LabelMap -> [HackInstruction] -> Either T.Text [Word16]
compile labels vars lines = forM lines compileInstruction
  where
    compileInstruction (A (Imm i)) =
      if i > 32767 || i < -32768
        then Left "Invalid immediate value"
        else Right $ fromIntegral i
    compileInstruction (A (Label l)) =
      case M.lookup l labels of
        Just offset -> Right $ fromIntegral offset
        Nothing -> case M.lookup l vars of
          Just offset -> Right $ fromIntegral offset
          Nothing     -> Left $ "Invalid label: " <> l
    compileInstruction (C dest src jmp) = Right $
      0xe000 .|. compileSourceOperand src  .|. compileDestOperand dest .|. compileJump jmp

    compileSourceOperand :: Source -> Word16
    compileSourceOperand src = compileSourceOperand' src `shift` 6

    compileSourceOperand' S0                = 0x2a
    compileSourceOperand' S1                = 0x3f
    compileSourceOperand' SNeg1             = 0x3a
    compileSourceOperand' (SReg RegA)       = 0x30
    compileSourceOperand' (SReg RegM)       = 0x70
    compileSourceOperand' (SReg RegD)       = 0x0c
    compileSourceOperand' (SNot RegA)       = 0x31
    compileSourceOperand' (SNot RegM)       = 0x71
    compileSourceOperand' (SNot RegD)       = 0x0d
    compileSourceOperand' (SNeg RegA)       = 0x33
    compileSourceOperand' (SNeg RegM)       = 0x73
    compileSourceOperand' (SNeg RegD)       = 0x0f
    compileSourceOperand' (SRegPlus1 RegA)  = 0x37
    compileSourceOperand' (SRegPlus1 RegM)  = 0x77
    compileSourceOperand' (SRegPlus1 RegD)  = 0x1f
    compileSourceOperand' (SRegMinus1 RegA) = 0x32
    compileSourceOperand' (SRegMinus1 RegM) = 0x72
    compileSourceOperand' (SRegMinus1 RegD) = 0x0e
    compileSourceOperand' SDPlusA           = 0x02
    compileSourceOperand' SDPlusM           = 0x42
    compileSourceOperand' SDMinusA          = 0x13
    compileSourceOperand' SDMinusM          = 0x53
    compileSourceOperand' SAMinusD          = 0x07
    compileSourceOperand' SMMinusD          = 0x47
    compileSourceOperand' SDAndA            = 0x00
    compileSourceOperand' SDAndM            = 0x40
    compileSourceOperand' SDOrA             = 0x15
    compileSourceOperand' SDOrM             = 0x55

    compileDestOperand dests = compileDestOperand' dests `shift` 3
    compileDestOperand' :: [Reg] -> Word16
    compileDestOperand' = foldl' compileDestBit (0 :: Word16)

    compileDestBit :: Word16 -> Reg -> Word16
    compileDestBit b RegA = b .|. 4
    compileDestBit b RegM = b .|. 1
    compileDestBit b RegD = b .|. 2

    compileJump JNone = 0
    compileJump JGt   = 1
    compileJump JEq   = 2
    compileJump JGe   = 3
    compileJump JLt   = 4
    compileJump JNe   = 5
    compileJump JLe   = 6
    compileJump Jmp   = 7