module Compiler where

import Parser
import Control.Monad (void) 
import System.IO
import Control.Monad.Trans.State
import Control.Monad.IO.Class (liftIO)
import qualified Data.Map as M

forthPrelude :: Handle -> IO () 
forthPrelude h = do
    hPutStr h
             "global _start\n\
             \print_number:\n\
             \    push rbp\n\
             \    mov rbp, rsp\n\
             \    sub rsp, 128 \n\
             \    mov rdx, -1 \n\
             \    jmp l2\n\
             \    add rsp, 28\n\
             \    pop rbp\n\
             \    ret\n\
             \l1:\n\
             \    dec  rdx\n\
             \    imul rax, rbx, 1717986919\n\
             \    shr  rax, 34\n\
             \    imul rcx, rax, 10\n\
             \    sub  rbx, rcx\n\
             \    add  rbx, '0'\n\
             \    mov  qword[rbp+8*rdx], rbx\n\
             \\n\
             \    mov rbx, rax\n\
             \l2:\n\
             \    cmp rbx, 0\n\
             \    jne l1\n\
             \    mov qword[rbp-8], `\\n` \n\
             \    mov rax, 1\n\
             \    mov rdi, 1\n\
             \    lea rsi, [rbp+8*rdx]\n\
             \    neg rdx\n\
             \    imul rdx, 8\n\
             \    syscall\n\
             \    leave\n\
             \    ret\n\
             \_start:\n"

data Environment = Environment { label :: Int
                               , environment :: M.Map String ZorthAST
                               }

handleSymbol :: Handle -> ZorthExpr -> StateT Environment IO ()

handleSymbol _ (ZorthASTComment _) = return ()

handleSymbol h (ZorthASTInteger i) = 
    do
        liftIO $ hPutStrLn h $ "    push "<>show i
        return ()

handleSymbol h (ZorthASTWord "+") = do
    liftIO $ hPutStr h 
             "    pop rbx\n\
             \    pop rax\n\
             \    add rax, rbx\n\
             \    push rax\n"
    return ()

handleSymbol h (ZorthASTWord "-") = do
    liftIO $ hPutStr h 
             "    pop rbx\n\
             \    pop rax\n\
             \    sub rax, rbx\n\
             \    push rax\n"
    return ()

handleSymbol h (ZorthASTWord "*") = do
    liftIO $ hPutStr h 
             "    pop rbx\n\
             \    pop rax\n\
             \    imul rax, rbx\n\
             \    push rax\n"
    return ()

handleSymbol h (ZorthASTWord "/") = do
    liftIO $ hPutStr h 
             "    pop rbx\n\
             \    pop rax\n\
             \    cqo\n\
             \    idiv rbx\n\
             \    push rax\n"
    return ()

handleSymbol h (ZorthASTWord "dup") = do
    liftIO $ hPutStr h 
             "    pop rax\n\
             \    push rax\n\
             \    push rax\n"
    return ()

handleSymbol h (ZorthASTWord "display") = do
    liftIO $ hPutStr h 
             "    push rsp\n\
             \    mov rbp, rsp\n\
             \    mov rax, 1\n\
             \    mov rdi, 1\n\
             \    mov rsi, [rbp]\n\
             \    mov rdx, 1\n\
             \    syscall\n\
             \    pop rsp\n\
             \    add rsp, 8\n"
    return ()

handleSymbol h (ZorthASTWord "mem") = do
    liftIO $ hPutStr h 
             "    push mem\n"
    return ()

handleSymbol h (ZorthASTWord "!") = do
    liftIO $ hPutStr h
             "    pop rbx\n\
             \    pop rax\n\
             \    mov qword[rax], rbx\n"
    return ()

handleSymbol h (ZorthASTWord "@") = do
    liftIO $ hPutStr h
             "    pop  rax\n\
             \    push qword[rax]\n"
    return ()

handleSymbol h (ZorthASTWord "rot") = do
    liftIO $ hPutStr h 
             "    pop rax\n\
             \    pop rbx\n\
             \    pop rcx\n\
             \    push rbx\n\
             \    push rax\n\
             \    push rcx\n"
    return ()

handleSymbol h (ZorthASTWord "swap") = do
    liftIO $ hPutStr h 
             "    pop rax\n\
             \    pop rbx\n\
             \    push rax\n\
             \    push rbx\n"
    return ()

handleSymbol h (ZorthASTWord "drop") = do
    liftIO $ hPutStr h "    add rsp, 8\n"
    return ()

handleSymbol h (ZorthASTWord ".") = do
    liftIO $ hPutStr h 
             "    pop rbx\n\
             \    call print_number\n"
    return ()

handleSymbol h (ZorthASTWord "=") = do
    liftIO $ truthOperator h "sete"
    return ()

handleSymbol h (ZorthASTWord "/=") = do
    liftIO $ truthOperator h "setne"
    return ()

handleSymbol h (ZorthASTWord ">") = do
    liftIO $ truthOperator h "setg"
    return ()

handleSymbol h (ZorthASTWord "<") = do
    liftIO $ truthOperator h "setl"
    return ()

handleSymbol h (ZorthASTWord ">=") = do
    liftIO $ truthOperator h "setge"
    return ()

handleSymbol h (ZorthASTWord "<=") = do
    liftIO $ truthOperator h "setle"
    return ()

handleSymbol h (ZorthASTWord w) = do
    s <- get
    _ <- compileZorthASTState h (environment s M.! w) -- parent env to child and discard
    return ()

handleSymbol _ (ZorthASTWordDecl (name,ast)) = do
    (Environment labels s) <- get
    put $ Environment labels (M.insert name ast s) -- maybe use `lens`?..
    return ()

-- honestly, this feels kinda janky, the way I manipulate state
handleSymbol h (ZorthASTIfElse (ifBranch,elseBranch)) = do
    (Environment l s) <- get
    let l1 = show $ l+1
    let l2 = show $ l+2
    put $ Environment (l+2) s
    liftIO $ hPutStr h $
             "    pop rax\n\
             \    cmp rax, 0\n\
             \    je .L"<>l1<>"\n"
    compileZorthASTState h ifBranch
    (Environment l' _) <- get
    put $ Environment l' s
    liftIO $ hPutStr h $ "    jmp .L"<>l2<>"\n.L"<>l1<>":\n"
    compileZorthASTState h elseBranch
    (Environment l'' _) <- get
    put $ Environment l'' s
    liftIO $ hPutStr h $ ".L"<>l2<>":\n"

handleSymbol h (ZorthASTWhile (condition,body)) = do
    (Environment l s) <- get
    let bodyl = show $ l+1
    let conditionl = show $ l+2
    let restl = show $ l+3
    put $ Environment (l+3) s
    liftIO $ hPutStrLn h $ "    jmp .L"<>conditionl
    liftIO $ hPutStrLn h $ ".L"<>bodyl<>":"
    compileZorthASTState h body 
    (Environment l' _) <- get
    put $ Environment l' s
    liftIO $ hPutStrLn h $ "\n.L"<>conditionl<>":"
    compileZorthASTState h condition
    (Environment l'' _) <- get
    put $ Environment l'' s
    liftIO $ hPutStrLn h "    pop rax"
    liftIO $ hPutStrLn h "    cmp rax,0"
    liftIO $ hPutStrLn h $ "    jg .L"<>bodyl
    liftIO $ hPutStrLn h $ ".L"<>restl<>":"

truthOperator :: Handle -> String -> IO ()
truthOperator h s =
    hPutStr h $
    "    pop rbx\n\
    \    pop rax\n\
    \    cmp rax, rbx\n\
    \    "<>s<>" al\n\
    \    movsx rax, al\n\
    \    push rax\n"

compileZorthASTState :: Handle -> ZorthAST -> StateT Environment IO ()
compileZorthASTState h = foldl (\m x -> m >> handleSymbol h x) (return ())

compileZorthAST :: Handle -> ZorthAST -> Environment -> IO ()
compileZorthAST h ast s = void $ runStateT (compileZorthASTState h ast) s

compileZorth :: Handle -> ZorthAST -> IO ()
compileZorth _ [] = return ()
compileZorth h xs = do
    forthPrelude h
    compileZorthAST h xs $ Environment 0 M.empty
    hPutStr h "    mov rax,60\n\
              \    mov rdi,0\n\
              \    syscall\n"
    hPutStr h "section .bss\n\
              \    mem: resq 640000\n"