{-|
Module      : Pulsar.Decode
Description : Decode data into DCPU16 assembly.
Copyright   : (c) Jesse Haber-Kucharsky, 2014
License     : Apache-2.0
-}

{-# LANGUAGE OverloadedStrings #-}

module Pulsar.Decode
  (file
  ,fileRaw
  ,program
  ,DecoderError (..))
  where

import           Pulsar.Ast             hiding (ByteString)
import           Pulsar.Output.Friendly
import           Pulsar.Spec            (SystemSpec)
import qualified Pulsar.Spec            as Spec
import           Pulsar.SymbolTable     (SymbolTable)
import qualified Pulsar.SymbolTable     as SymbolTable

import           Control.Monad          (forM)
import           Data.Binary.Get
import           Data.Bits
import qualified Data.ByteString        as S
import qualified Data.ByteString.Char8  as S.Char8
import           Data.ByteString.Lazy   (ByteString)
import qualified Data.ByteString.Lazy   as LBS
import           Data.Foldable          (toList)
import           Data.List              (groupBy, intersect, nub, sort)
import           Data.Map               (Map)
import qualified Data.Map               as Map
import           Data.Maybe             (fromMaybe, mapMaybe)
import           Data.Monoid            ((<>))
import qualified Data.Text              as Text
import           Data.Vector.Unboxed    (Vector)
import qualified Data.Vector.Unboxed    as Vec

data DecoderError = DecoderError FilePath deriving (Eq, Show)

instance Friendly DecoderError where
  friendly (DecoderError path) =
    "Unable to decode \"" <> Text.pack path <> "\"."

-- | Decode a whole file on disk.
file :: SystemSpec -> FilePath -> IO (Maybe (Ast Word16))
file spec path = LBS.readFile path >>= program spec

fileRaw :: FilePath -> IO (Maybe (Vector Word16))
fileRaw path =
  do bs <- LBS.readFile path
     case runGetOrFail (unfoldGet getNext) bs of
       Left _ -> return Nothing
       Right (_, _, result) -> return . Just . Vec.fromList $ result

unfoldGet :: Get (Maybe a) -> Get [a]
unfoldGet g =
  do x <- g
     case x of
       Nothing -> return []
       Just v ->
         do vs <- unfoldGet g
            return $ v : vs

-- | Decode a ByteString.
program :: SystemSpec
        -> ByteString
        -> IO (Maybe (Ast Word16))
program spec bs =
  -- Start out by parsing the symbol table.
  let (bs', headerSizeWords, symbols, symbolsStartLoc) =
        case runGetOrFail (symbolTable spec) bs of
          Left _ -> failCase
          Right (remainingBs, numBytes, symTabResult) ->
            case symTabResult of
              Nothing -> failCase
              Just (symTab, symTabLoc) -> (remainingBs, fromIntegral numBytes `div` 2, symTab, symTabLoc)
        where
          failCase = (bs, 0, SymbolTable.empty, 0)

      programBytes = LBS.take
                       (fromIntegral $ locToByteOffsetFromHeader symbolsStartLoc)
                       bs'
  in
  do result <- bareProgram spec programBytes
     case result of
       Nothing -> return Nothing
       Just ast ->
             -- Correct the offsets in the AST to account for the size of the
             -- symbol table header.
         let ast' = map (fmap (+ headerSizeWords)) ast
             statementLocs = concatMap toList ast'
             -- The only valid locations for a label are those where a
             -- statement begins.
             validLabelLocs = labelLocs ast' `intersect` statementLocs
             -- The default labels, without an symbol table knowledge.
             rawLabelLocsTab = Map.fromList $ zip validLabelLocs
                                                  labelNames
             labelLocsTab = insertMany (map swap . SymbolTable.toList $ symbols)
                                       rawLabelLocsTab

         in return $ Just
                   . optimize
                   . replaceJumps labelLocsTab
                   . insertLabels labelLocsTab $ ast'

swap :: (a, b) -> (b, a)
swap (x, y) = (y, x)

insertMany :: Ord k => [(k, a)] -> Map k a -> Map k a
insertMany pairs x = foldr (\(k, v) m -> Map.insert k v m) x pairs

-- | Condense statements that can be combined.
optimize :: Ast Word16 -> Ast Word16
optimize = concatMap combine . groupBy isMatchingDat
  where
    -- | Multiple .dat directives of the same value can be replaced with a .fill.
    combine xs@(Statement loc (Dat [expr]) : _)
      | length xs >= 3 =
        [Statement loc (Fill (Expression loc (Word (fromIntegral $ length xs))) expr)]
    combine xs = xs

    isMatchingDat (Statement _ (Dat [Expression _ (Word x)]))
                  (Statement _ (Dat [Expression _ (Word y)])) = x == y
    isMatchingDat _ _ = False

labelNames :: [S.ByteString]
labelNames = [S.Char8.pack $ "label" ++ show index | index <- [1 :: Integer ..]]

-- | Insert label definitions based on where labels are referenced in the
-- decoded program by key instructions.
insertLabels :: Map Word16 S.ByteString
             -> Ast Word16
             -> Ast Word16
insertLabels labelLocsTab = concatMap insert
  where
    insert :: Statement Word16 -> Ast Word16
    insert s@(Statement loc _) =
      case Map.lookup loc labelLocsTab of
        Nothing -> [s]
        Just labelName -> [Statement loc (Label labelName), s]

-- | Replace a jump to an address by a jump to a named label.
replaceJumps :: Map Word16 S.ByteString -> Ast Word16 -> Ast Word16
replaceJumps tab = map replace
  where
    replace (Statement pos s) = Statement pos $
      case strip s of
        Ins (Op () "SET")
            (Address () Pc)
            (Address () (Dir (Expression () (Word loc)))) ->
              tag pos $ Ins (Op () "SET")
                            (Address () Pc)
                            (Address () (Dir (find loc)))
        SpecialIns (SpecialOp () "JSR")
                   (Address () (Dir (Expression () (Word loc)))) ->
              tag pos $ SpecialIns (SpecialOp () "JSR")
                                   (Address () (Dir (find loc)))
        _ -> s

    find :: Word16 -> Expression ()
    find loc = Expression () $
      case Map.lookup loc tab of
        Just name -> Ident name
        Nothing -> Word loc

-- | Offsets into the data image where labels could be located.
labelLocs :: Ast Word16 -> [Word16]
labelLocs = sort . nub . mapMaybe (get . strip)
  where
    get (Statement _ (Ins (Op _ "SET")
                          (Address () Pc)
                          (Address () (Dir (Expression () (Word loc)))))) = Just loc
    get (Statement _ (SpecialIns (SpecialOp _ "JSR")
                                 (Address () (Dir (Expression () (Word loc)))))) = Just loc
    get _ = Nothing

-- | The encoded program, without any symbol table.
bareProgram :: SystemSpec -> ByteString -> IO (Maybe (Ast Word16))
bareProgram spec initial =
  do result <- loop 0 initial
     case result of
       [] -> return Nothing
       ast -> return . Just $ ast
  where
    loop :: Word16
         -> ByteString
         -> IO (Ast Word16)
    loop loc bs =
      if LBS.null bs
         then return []
         else
           case runGetOrFail (getStatement spec) bs of
             Left _ -> return []
             Right (bs', bytes, st) ->
               let numWordsRead = fromIntegral $ bytes `div` 2
                   loc' = loc + numWordsRead
               in
                 do next <- loop loc' bs'
                    return $ tag loc st : next

-- | If the SymbolTable is successfully decoded, then the second element of the
-- tuple is the location in the data image (in DCPU16 words) where the table
-- resides.
symbolTable :: SystemSpec -> Get (Maybe (SymbolTable, Word16))
symbolTable spec =
  do (Statement () first) <- getStatement spec
     case first of
       Ins (Op () "SET")
           (Address () Pc)
           (Address () (Dir (Expression () (Word 3)))) ->
         do hasSig <- getSignature
            if not hasSig
               then return Nothing
               else
                 do tableLoc <- getWord
                    lookAhead $
                      do skip (locToByteOffsetFromHeader tableLoc)
                         numEntries <- getWord
                         entries <- forM [1 .. numEntries] (const getEntry)
                         return $ Just (SymbolTable.fromList entries, tableLoc)
       _ -> return Nothing
  where
    getSignature :: Get Bool
    getSignature =
      do isEmpty_ <- isEmpty
         if isEmpty_
            then return False
            else
              do w <- getWord
                 return (w == SymbolTable.signature)

    getEntry :: Get (S.ByteString, Word16)
    getEntry =
      do len <- fmap fromIntegral getWord
         name <- getByteString len
         _ <- getByteString (if odd len then 1 else 0)
         loc <- getWord
         return (name, loc)

-- | This function converts the index of the 16 bit word indicating the location
-- of the symbol table in the image to a byte offset from the end of the symbol
-- table header.
locToByteOffsetFromHeader :: Word16 -> Int
locToByteOffsetFromHeader n =
  -- To prevent overflow.
  let n' = fromIntegral n :: Word32
  in fromIntegral $ (2 * n') - (3 * 2)

getStatement :: SystemSpec -> Get (Statement ())
getStatement spec =
  do w <- getWord
     insM <- choice [getIns spec w
                    ,getSpecialIns spec w]
     return $ fromMaybe (Statement () (Dat [Expression () (Word w)])) insM

getIns :: SystemSpec -> Word16 -> Get (Maybe (Statement ()))
getIns spec w =
  let opM = opCode spec $ w .&. 0x1f
  in
    do addrAM <- getAddress spec ContextA $ (w .&. 0xfc00) `shiftR` 10
       addrBM <- getAddress spec ContextB $ (w .&. 0x3e0) `shiftR` 5
       return $
         do op <- opM
            addrA <- addrAM
            addrB <- addrBM
            return $ Statement () (Ins op addrB addrA)

opCode :: SystemSpec -> Word16 -> Maybe (Op ())
opCode = Spec.decode . Spec.opCodes

getSpecialIns :: SystemSpec
              -> Word16
              -> Get (Maybe (Statement ()))
getSpecialIns spec w =
  let prefixM = zeroPrefix $ w .&. 0x10
      opM = specialOpCode spec $ (w .&. 0x3e0) `shiftR` 5
  in
    do addrM <- getAddress spec ContextA $ (w .&. 0xfc00) `shiftR` 10
       return $
        do _ <- prefixM
           op <- opM
           addr <- addrM
           return $ Statement () (SpecialIns op addr)
  where
    zeroPrefix 0 = return $ Just ()
    zeroPrefix _ = return Nothing

specialOpCode :: SystemSpec -> Word16 -> Maybe (SpecialOp ())
specialOpCode = Spec.decode . Spec.specialOpCodes

getAddress :: SystemSpec
           -> AddressContext
           -> Word16
           -> Get (Maybe (Address ()))
getAddress spec context w =
  choice [getRegDir spec w
         ,getRegInd spec w
         ,getRegIndOff spec w
         ,getFastDir w
         ,getDir w
         ,getInd w
         ,getPeek w
         ,getPick w
         ,getPushOrPop context w
         ,getSp w
         ,getPc w
         ,getEx w]

getPushOrPop :: AddressContext -> Word16 -> Get (Maybe (Address ()))
getPushOrPop context 0x18 =
  return . fmap (Address ()) $ case context of
                                 ContextA -> Just Pop
                                 ContextB -> Just Push
getPushOrPop _ _ = return Nothing

getFastDir :: Word16 -> Get (Maybe (Address ()))
getFastDir 0x20 = return . Just . Address () . Dir . Expression () . Word $ 0xffff
getFastDir w
  | w >= 0x21 && w <= 0x3f =
    return . Just . Address () . Dir . Expression () . Word $ w - 0x21
  | otherwise = return Nothing

getDir :: Word16 -> Get (Maybe (Address ()))
getDir 0x1f = fmap (Address () . Dir) <$> getNextExpr
getDir _ = return Nothing

getPeek :: Word16 -> Get (Maybe (Address ()))
getPeek 0x19 = return . Just $ Address () Peek
getPeek _ = return Nothing

getPick :: Word16 -> Get (Maybe (Address ()))
getPick 0x1a = fmap (Address () . Pick) <$> getNextExpr
getPick _ = return Nothing

getSp :: Word16 -> Get (Maybe (Address ()))
getSp 0x1b = return . Just $ Address () Sp
getSp _ = return Nothing

getPc :: Word16 -> Get (Maybe (Address ()))
getPc 0x1c = return . Just $ Address () Pc
getPc _ = return Nothing

getEx :: Word16 -> Get (Maybe (Address ()))
getEx 0x1d = return . Just $ Address () Ex
getEx _ = return Nothing

getInd :: Word16 -> Get (Maybe (Address ()))
getInd 0x1e = fmap (Address () . Ind) <$> getNextExpr
getInd _ = return Nothing

getRegDir :: SystemSpec -> Word16 -> Get (Maybe (Address ()))
getRegDir spec = return . fmap (Address () . RegDir) . register spec

getRegInd :: SystemSpec -> Word16 -> Get (Maybe (Address ()))
getRegInd spec =
  return . fmap (Address () . RegInd) . register spec . subtract 8

getRegIndOff :: SystemSpec -> Word16 -> Get (Maybe (Address ()))
getRegIndOff spec w =
  case register spec (w - 0x10) of
    Nothing -> return Nothing
    Just r -> fmap (Address () . RegIndOff r) <$> getNextExpr

register :: SystemSpec -> Word16 -> Maybe (Register ())
register = Spec.decode . Spec.registers

-- | Decode one of many possibilities, with backtracking.
choice :: [Get (Maybe a)] -> Get (Maybe a)
choice [] = return Nothing
choice (x : xs) =
  do result <- lookAheadM x

     case result of
       Just _ -> return result
       Nothing -> choice xs

getNextExpr :: Get (Maybe (Expression ()))
getNextExpr = fmap (Expression () . Word) <$> getNext

getNext :: Get (Maybe Word16)
getNext =
  do isEmpty_ <- isEmpty
     if isEmpty_
        then return Nothing
        else Just <$> getWord

getWord :: Get Word16
getWord = getWord16be
