{-# LANGUAGE DeriveFunctor, GeneralizedNewtypeDeriving #-}

module Compiler(compile) where

import Control.Monad.State.Strict
import Control.Monad.Writer.Strict
import Data.List
import qualified System.Info as System (os)

import AST


-- Execution starts in the middle of the tape
tapeLength :: Offset
tapeLength = 60000


newtype LineWriter a = LineWriter {unLineWriter :: WriterT [String] (State Int) a}
  deriving (Functor, Applicative, Monad, MonadWriter [String], MonadState Int)

runLineWriter :: LineWriter a -> String
runLineWriter = (++ "\n") . intercalate "\n" . flip evalState 1 . execWriterT . unLineWriter

emit0 :: String -> LineWriter ()
emit0 = tell . pure

emit :: String -> LineWriter ()
emit = tell . pure . ('\t' :)

genId :: LineWriter Int
genId = state $ \i -> (i, i + 1)


mangleName :: String -> String
mangleName f = case System.os of
    "darwin" -> '_' : f
    "linux" -> f
    _ -> undefined

generateCall :: String -> String
generateCall f = case System.os of
    "darwin" -> "call " ++ mangleName f
    "linux" -> "call [" ++ mangleName f ++ " wrt ..got]"
    _ -> undefined


compile :: Program -> String
compile (Program inss) = runLineWriter $ do
    emit0 prologue
    mapM_ compileIns inss
    emit0 epilogue

cursorReg :: String
cursorReg = "rbx"

cursorOffset :: Offset -> String
cursorOffset 0 = cursorReg
cursorOffset off = cursorReg ++ " + " ++ show off

compileIns :: Instruction -> LineWriter ()
compileIns (IAdd v off) = emit $ "add byte [" ++ cursorOffset off ++ "], " ++ show v
compileIns (ISet v off) = emit $ "mov byte [" ++ cursorOffset off ++ "], " ++ show v
compileIns (ICopy _ _ 0) = return ()
compileIns (ICopy from to m) = do
    emit $ "mov al, [" ++ cursorOffset from ++ "]"
    case m of
        1 -> emit $ "add [" ++ cursorOffset to ++ "], al"
        -1 -> emit $ "sub [" ++ cursorOffset to ++ "], al"
        _ | Just p <- isTwoPower m -> do
                emit $ "shl al, " ++ show p
                emit $ "add [" ++ cursorOffset to ++ "], al"
          | Just p <- isTwoPower (-m) -> do
                emit $ "shl al, " ++ show p
                emit $ "sub [" ++ cursorOffset to ++ "], al"
          | otherwise -> do
                emit $ "mov cl, " ++ show m
                emit "mul byte cl"
                emit $ "add [" ++ cursorOffset to ++ "], al"
compileIns (ISlide off) = emit $ "add " ++ cursorReg ++ ", " ++ show off
compileIns (ILoop inss off) = do
    loopid <- genId
    emit $ "cmp byte [" ++ cursorOffset off ++ "], 0"
    emit $ "jz .Laf_" ++ show loopid
    emit0 $ ".Lbd_" ++ show loopid ++ ":"
    mapM_ compileIns inss
    emit $ "cmp byte [" ++ cursorOffset off ++ "], 0"
    emit $ "jnz .Lbd_" ++ show loopid
    emit0 $ ".Laf_" ++ show loopid ++ ":"
compileIns (IInput off) = do
    emit $ generateCall "getchar"
    emit $ "mov [" ++ cursorOffset off ++ "], al"
compileIns (IOutput off) = do
    emit $ "movzx edi, byte [" ++ cursorOffset off ++ "]"
    emit $ generateCall "putchar"
compileIns IStart = return ()

isTwoPower :: Byte -> Maybe Int
isTwoPower v = findIndex (==v) (take 8 $ iterate (* 2) 1)


prologue :: String
prologue =
         "global " ++ mangleName "main"
    ++ "\nextern " ++ intercalate ", "
            (map mangleName ["calloc", "free", "putchar", "getchar", "write"])
    ++ "\ndefault rel"
    ++ "\nsection .text"

    ++ "\n" ++ mangleName "main" ++ ":"
    ++ "\n\tpush rbp"
    ++ "\n\tmov rbp, rsp"
    ++ "\n\tpush rbx"
    ++ "\n\tsub rsp, 8"
    ++ "\n\tmov edi, " ++ show tapeLength
    ++ "\n\tmov esi, 1"
    ++ "\n\t" ++ generateCall "calloc"
    ++ "\n\ttest rax, rax"
    ++ "\n\tjz .allocation_failure"
    ++ "\n\tlea rbx, [rax + " ++ show (tapeLength `div` 2) ++ "]  ; rbx = cursor"
    ++ "\n\tmov [rsp], rax  ; [rsp] = start of tape buffer"
    ++ "\n"

epilogue :: String
epilogue =
       "\n\tmov rdi, [rsp]"
    ++ "\n\t" ++ generateCall "free"
    ++ "\n\txor eax, eax"
    ++ "\n.return:"
    ++ "\n\tadd rsp, 8"
    ++ "\n\tpop rbx"
    ++ "\n\tmov rsp, rbp"
    ++ "\n\tpop rbp"
    ++ "\n\tret"
    ++ "\n.allocation_failure:"
    ++ "\n\tsub rsp, 32"
    ++ "\n\tmov rax, 0x697461636f6c6c41  ; \"Allocati\""
    ++ "\n\tmov rbx, 0x756c696166206e6f  ; \"on failu\""
    ++ "\n\tmov ecx, 0x000a6572          ; \"re\\n\\0\""
    ++ "\n\tmov [rsp], rax"
    ++ "\n\tmov [rsp+8], rbx"
    ++ "\n\tmov [rsp+16], ecx"
    ++ "\n\tmov edi, 2"
    ++ "\n\tmov rsi, rsp"
    ++ "\n\tmov edx, 20"
    ++ "\n\t" ++ generateCall "write"
    ++ "\n\tadd rsp, 32"
    ++ "\n\tmov eax, 1"
    ++ "\n\tjmp .return"