module Language.HAsm.Parse where

import Data.Char
import Data.Either (lefts, rights)
import Data.List (intercalate, nub, sort)
import qualified Data.Map as M

import qualified Control.Arrow as Arr
import Control.Applicative ((<$>), (<*), liftA2)
import Control.Monad (when)

import Numeric

import Text.Parsec
import Text.Parsec.String (Parser)

import Language.HAsm.Types
import Language.HAsm.X86.CPU

hasmParseFile :: FilePath -> IO (Either ParseError ParseResult)
hasmParseFile fname = hasmParseWithSource fname <$> readFile fname

hasmParseWithSource :: FilePath -> String -> Either ParseError ParseResult
hasmParseWithSource fname ss = parse (asmfile <* eof) fname ss

asmfile :: Parser ParseResult
asmfile = do
  blanklines
  --mbasmspaces
  many $ do
    pos <- getPosition
    stmt <- asmstmt
    blanklines
    --mbasmspaces
    return (stmt, toHasmPos pos)

asmstmt :: Parser HasmStatement
asmstmt = try asmlabel <|> try asmdir <|> asmop <?> "assembly statement"

blanklines = many ((linecomment <|> newline) >> mbasmspaces) >> return ()
linecomment = do
  char '#'; many (noneOf "\n"); newline

dirchar = letter <|> char '_'
idschar = letter <|> oneOf "._$"
idchar = idschar <|> digit

symbol = cons idschar (many idchar)

endstmt = mbasmspaces >> optional ((oneOf ";" >> return ()) <|> eof)
endlbl = (mbasmspaces >> blanklines) <|> eof

mbasmspaces = many (oneOf " \t")

cons a b = liftA2 (:) a b
opchar = letter <|> digit
opcode = cons letter (many1 opchar)
regname = cons letter (many1 opchar)

asmdir = do
  char '.'
  dir <- many1 dirchar
  mbasmspaces
  dirargs <- sepBy (many1 $ noneOf "\n,") (char ',' >> mbasmspaces)
  endstmt
  case readDirective dir dirargs of
    Right d -> return $ HasmStDirective d
    Left e -> parserFail e

asmlabel = HasmStLabel <$> (symbol <* char ':' <* endlbl) <?> "label"

asmop = do
  op <- opcode
  mbasmspaces
  opnds <- sepBy asmopnd (char ',' >> mbasmspaces)
  endstmt
  case readOperation op opnds of
    Left e -> parserFail e
    Right o -> return $ HasmStInstr [] o

asmopnd :: Parser OpOperand
asmopnd = try (OpndImm <$> asmimm) <|>
          try (OpndReg <$> asmregister) <|>
          try asmmem <?> "opcode operand"

immToDispl iv@(ImmL imm)
    | isWord8 iv = Displ8 $ fromIntegral imm
    | otherwise = Displ32 imm

asmdispl = try (immToDispl <$> readIntegerLit)
           <|> (DisplLabel <$> symbol)
           <?> "displacement before ("

asmmem = do
  mbDspl <- optionMaybe asmdispl
  mbSIB <- optionMaybe asmsib
  when (mbDspl == Nothing && mbSIB == Nothing) $ fail "no displacement and no SIB"
  let dspl = maybe NoDispl id mbDspl
  let sib = maybe noSIB id mbSIB
  return $ OpndRM sib dspl


{-- SIBs and memory refs --}
data SIBPart = SIBReg GPRegister | SIBScale Word8 | SIBNothing
immToSIBScale (ImmL imm) = SIBScale $ fromIntegral imm

sibpart = try (SIBReg <$> asmgpregister) <|>
          try (immToSIBScale <$> readIntegerLit) <|>
          return SIBNothing

sibparts = between (char '(') (char ')') $ sepBy sibpart (char ',' >> mbasmspaces)

asmsib = do
  sps <- sibparts
  case sps of
    -- (%reg)
    [SIBReg base] -> return $ SIB 1 Nothing (Just base)
    -- (%reg, %reg)
    [SIBReg base, SIBReg ind] -> return $ SIB 1 (Just ind) (Just base)
    -- (%reg, %reg, )
    [SIBReg base, SIBReg ind, SIBNothing] -> return $ SIB 1 (Just ind) (Just base)
    -- (, %reg)
    [SIBNothing, SIBReg ind] -> return $ SIB 1 (Just ind) Nothing
    -- (%reg, scale)
    [SIBReg ind, SIBScale scale ] -> makeSIB scale (Just ind) Nothing
    -- (, %reg, scale)
    [SIBNothing, SIBReg ind, SIBScale scale ] -> makeSIB scale (Just ind) Nothing
    -- (%reg, %reg, scale)
    [SIBReg base, SIBReg ind, SIBScale scale ] -> makeSIB scale (Just ind) (Just base)
    -- everything else:
    _ -> fail "SIB format is invalid"
  where
    makeSIB scale mbInd mbBase =
      if scale `elem` [1,2,4,8]
      then return $ SIB (fromIntegral scale) mbInd mbBase
      else fail "Scale is not 1,2,4,8"

asmregister = do
  char '%'
  reg <- regname
  case mbRegByName reg of
    Just r -> return r
    _ -> parserFail $ "Not a register name: " ++ reg

asmgpregister = do
  RegL reg <- asmregister
  return reg

{-- Number literals --}
asmimm = do
  char '$'
  try readIntegerLit <|> (ImmS <$> symbol) <?> "integer or symbol in immediate value"

readIntegerLit = (ImmL . fromIntegral) <$> (intneg <|> intparse <?> "number")

intneg = do
    char '-'
    num <- intparse
    return $ negate num

hexint = do
  ds <- many1 (oneOf "0123456789abcdefABCDEF")
  case readHex ds of
    [(num, "")] -> return num
    _ -> parserFail "failed to parse hexadecimal number"

octint = do
  ds <- many1 (oneOf "01234567")
  case readOct ds of
    [(num, "")] -> return num
    _ -> parserFail "failed to parse octal number"

intparse = do
  d <- digit
  if d == '0'
  then do
    try (char 'x' >> hexint) <|> try octint <|> return 0
  else do
    ds <- many digit
    case readDec (d:ds) of
      [(num, "")] -> return num
      _ -> parserFail "failed to parse a number"

{-- Directives --}
readDirective :: String -> [String] -> Either String Directive
readDirective dir args =
  case dir of
    "section" ->
      case args of
        [] -> Left $ ".section takes at least section name"
        [name] -> Right $ DirSection name 0 ""
        [name, subsect] -> case (reads :: ReadS Int) subsect of
                             [(nsubsect, "")] -> Right $ DirSection name nsubsect ""
                             _ -> Left $ "Failed to parse subsection: " ++ show subsect
        _ -> Left $ "invalid .section format"
    sect | sect `elem` ["text", "data"] ->
      case args of
        [] -> Right $ DirSection sect 0 ""
        [subsect] -> case (reads :: ReadS Int) subsect of
                        [(nsubsect, "")] -> Right $ DirSection sect nsubsect ""
                        _ -> Left $ "Failed to parse subsection: " ++ show subsect
    glbl | glbl `elem` ["global", "globl"] ->
      let eiSyms = map (parse (symbol <* eof) "") args
      in case lefts eiSyms of
          [] -> Right $ DirGlobal (rights eiSyms)
          (err:_) -> Left $ "failed to parse symbol " ++ show err
    "extern" ->
      let eiSyms = map (parse (symbol <* eof) "") args
      in case lefts eiSyms of
          [] -> Right $ DirExtern (rights eiSyms)
          (err:_) -> Left $ "failed to parse symbol " ++ show err
    "file" ->
      case args of
        [fname] -> Right $ DirFile fname
        _ -> Left ".file takes only one string parameter"
    "type" ->
      case args of
        [sym, ty] -> Right $ DirType sym ty
        _ -> Left ".type takes two arguments: .type <sym>, <type>"
    --"asciz" ->
    _ -> Left $ "Unknown directive: " ++ dir ++ " " ++ intercalate "," args


type OpSuffix = Char
type OpInfo = (Instr, [Int], [OpSuffix])

opsyntax :: [(String, OpInfo)]
opsyntax = [
  ("mov",   (OpMov,  [2],   "lwb")),
  ("add",   (OpAdd,  [2],   "lwb")),
  ("ret",   (OpRet,  [0,1], "w"  )),
  ("push",  (OpPush, [1],   "lwb")),
  ("pop",   (OpPop,  [1],   "lwb")),
  ("cmp",   (OpCmp,  [2],   "lwb")),
  ("int",   (OpInt,  [1],   "b"  )),
  ("jmp",   (OpJmp,  [1],   "lwb")),
  ("call",  (OpCall, [1],   "l"  )),
  ("imul",  (OpIMul, [1,2,3],"lwb")),

  ("je",    (OpJe,   [1],   ""   )),  -- ZF=1
  ("jz",    (OpJe,   [1],   ""   )),
  ("jne",   (OpJne,  [1],   ""   )),  -- ZF=0
  ("jnz",   (OpJne,  [1],   ""   )),
  ("jge",   (OpJge,  [1],   ""   )),  -- SF=OF
  ("jnl",   (OpJge,  [1],   ""   )),
  ("jl",    (OpJl,   [1],   ""   )),  -- SF<>OF
  ("jnge",  (OpJl,   [1],   ""   )),
  ("jle",   (OpJle,  [1],   ""   )),  -- ZF=1, SF<>OF
  ("jng",   (OpJle,  [1],   ""   )),
  ("jg",    (OpJg,   [1],   ""   )),  -- ZF=0, SF=OF
  ("jnle",  (OpJg,   [1],   ""   )),
  ("jnp",   (OpJnp,  [1],   ""   )),  -- PF=0
  ("jpo",   (OpJnp,  [1],   ""   )),
  ("jp",    (OpJp,   [1],   ""   )),  -- PF=1
  ("jpe",   (OpJp,   [1],   ""   )),
  ("jno",   (OpJno,  [1],   ""   )),  -- OF=0
  ("jo",    (OpJo,   [1],   ""   )),  -- OF=1
  ("jns",   (OpJns,  [1],   ""   )),  -- SF=0
  ("js",    (OpJs,   [1],   ""   )),  -- SF=1
  
  ("jae",   (OpJnc,  [1],   ""   )),  -- CF=0
  ("jnb",   (OpJnc,  [1],   ""   )),
  ("jnc",   (OpJnc,  [1],   ""   )),
  ("jc",    (OpJc,   [1],   ""   )),  -- CF=1
  ("jb",    (OpJc,   [1],   ""   )),
  ("jnae",  (OpJc,   [1],   ""   )),
  ("ja",    (OpJa,   [1],   ""   )),  -- CF=0, ZF=0
  ("jnbe",  (OpJa,   [1],   ""   )),
  ("jbe",   (OpJbe,  [1],   ""   )),  -- CF=1, ZF=1
  ("jna",   (OpJbe,  [1],   ""   )),

  ("jecxz", (OpJecxz,[1],   ""   ))  ]

opmap = M.fromList opsyntax

oplookup :: String -> Either String (OpInfo, OpSuffix)
oplookup opname =
  case M.lookup opname opmap of
    Just opinfo ->
      Right (opinfo, '?')
    Nothing ->
      -- search without possible suffixes:
      let suf = last opname
          opname' = init opname
      in case M.lookup opname' opmap of
        Just opinfo@(_, _, suffs) ->
          if suf `elem` suffs
          then Right (opinfo, suf)
          else Left $ "Invalid suffix for command " ++ opname'
        Nothing -> Left $ "Unknown instruction: " ++ opname

readOperation :: String -> [OpOperand] -> Either String Operation
readOperation opname opnds =
  case oplookup opname of
    Left e -> Left e
    Right (opinfo@(op, arglens, suffs), suf) ->
      if length opnds `elem` arglens
      then case unifyOperandTypes opinfo suf opnds of
             Right opnds' -> Right $ Operation op opnds'
             Left e -> Left $ "Operand type mismatch: " ++ e
      else Left $ "Opcode '" ++ opname ++ "' does not take " ++ show (length opnds) ++ " parameter(s)"

getOpndType :: OpOperand -> OpSuffix
getOpndType opnd =
  case opnd of
    OpndImm _        -> '?'    -- any, needs adjusting to OpndReg size
    OpndRM _ _       -> '?'    -- memory can be read by 1,2,4 bytes
    OpndReg (SReg _) -> 's'    -- compatible with w and l at the same time
    OpndReg (RegL _) -> 'l'    -- these registers are rigid
    OpndReg (RegW _) -> 'w'
    OpndReg (RegB _) -> 'b'

getOperandsTypes :: OpSuffix -> [OpOperand] -> [OpSuffix]
getOperandsTypes opsuf opnds =
  nub $ sort $ (opsuf : map getOpndType opnds)  -- what types are present?

unifyOperandTypes :: OpInfo -> OpSuffix -> [OpOperand] -> Either String [OpOperand]
unifyOperandTypes (opname, lens, sufs) opsuf opnds =
  let ts = filter (not . flip elem "s?") $ getOperandsTypes opsuf opnds
  in case ts of
      ""    -> Right opnds
      "l"   -> Right opnds  -- parser defaults to ImmL
      "w"   -> adjustImmediatesWith immToW opnds
      "b"   -> adjustImmediatesWith immToB opnds
      _     -> Left $ "Operand type mismatch"  -- TODO: e.g. movsx with different types

adjustImmediatesWith :: (ImmValue -> Either String ImmValue) ->
                        [OpOperand] -> Either String [OpOperand]
adjustImmediatesWith adjimm = mapM adj
  where
    adj (OpndImm imm) = OpndImm <$> adjimm imm
    adj op = Right op

toHasmPos :: SourcePos -> SrcPos
toHasmPos src = SrcPos (sourceName src) (sourceLine src) (sourceColumn src)
