{-# LANGUAGE OverloadedStrings #-}

module Language.Del.InRep where

import qualified Control.Monad.Trans.State as S
import qualified Data.List as L
import qualified Data.Text as T
import qualified Language.Del.AST as A
import           Language.Del.CompileDefs
import           Language.Del.TypeTable

data Atom
    = Name T.Text
    | InFloat Double
    | InInt Integer

data IExpr
    = Add IExpr IExpr
    | Mul IExpr IExpr
    | Sub IExpr IExpr
    | Div IExpr IExpr
    | Lt IExpr IExpr
    | Gt IExpr IExpr
    | LtE IExpr IExpr
    | GtE IExpr IExpr
    | Negate IExpr
    | Atomic Atom
    | Ternary IExpr IExpr IExpr
    | Builtin T.Text [IExpr]
    | Index IExpr Integer
    | Boolean Bool

data Stmt
    = Store T.Text Type IExpr
    | Return IExpr

data Program = Program [(T.Text, Type)] Type [Stmt]

toIR :: A.Def -> Artifact Program
toIR f@(A.Def _ r a _) = runCtx (gen f) >>= return . (Program a r)

genExpr :: A.Expr -> GenCtx IExpr

genExpr (A.Add xs) = nary Add A.Add xs
genExpr (A.Mul xs) = nary Mul A.Mul xs
genExpr (A.Sub x y) = binary Sub x y
genExpr (A.Div x y) = binary Div x y
genExpr (A.Lt x y) = binary Lt x y
genExpr (A.Gt x y) = binary Gt x y
genExpr (A.LtE x y) = binary LtE x y
genExpr (A.GtE x y) = binary GtE x y
genExpr (A.Negate x) = genExpr x >>= return . Negate
genExpr (A.Id n) = return $ Atomic $ Name n
genExpr e@(A.ScopedId _ _) = genExpr (A.canonicalId e)
genExpr (A.Number (A.DelInt i)) = return $ Atomic $ InInt i
genExpr (A.Number (A.DelFloat f)) = return $ Atomic $ InFloat f

genExpr (A.Pow x p) = do
    ex <- genExpr x
    ep <- genExpr p
    return $ Builtin "pow" [ex, ep]

genExpr (A.Apply (A.Id fun) args) = do
    atemps <- mapM genExpr args
    return $ Builtin fun atemps
genExpr (A.Apply _ _) = undefined

genExpr (A.If c t f) = do
    co <- genExpr c
    to <- genExpr t
    fo <- genExpr f
    return $ Ternary co to fo

genExpr (A.Subscript e s) = genExpr e >>= return . (flip Index) (subIdx s)
    where subIdx "x" = 0
          subIdx "y" = 1
          subIdx "z" = 2
          subIdx "0" = 0
          subIdx "1" = 1
          subIdx "2" = 2
          subIdx _ = undefined

genExpr (A.NativeCall t args) = genExpr (A.Apply (A.Id t) args)
genExpr (A.Boolean b) = return $ Boolean b

genExpr (A.Let _ _ _) = undefined

gen :: A.Def -> GenCtx ()
gen (A.Def _ _ _ b) = do
    baseref <- genExpr b
    _ <- nextTemp
    emit $ Return baseref

-- This is not the best name; really, this supports n-ary functions with n >= 2
nary :: (IExpr -> IExpr -> IExpr) -> ([A.Expr] -> A.Expr) -> [A.Expr] -> GenCtx IExpr
nary ictor ector (x:y:z:xs) = do
    xo <- genExpr x
    xso <- genExpr $ ector (y:z:xs)
    return (ictor xo xso)
nary ictor _ [x, y] = binary ictor x y
nary _ _ _ = undefined

binary :: (IExpr -> IExpr -> IExpr) -> A.Expr -> A.Expr -> GenCtx IExpr
binary ctor x y = do
    xo <- genExpr x
    yo <- genExpr y
    return $ ctor xo yo

data GenState
    = GenState [Stmt] Int -- list of statements, id of next temp var
type GenCtx a = S.StateT GenState Artifact a

-- Add an statement to the list of statements in the context
emit :: Stmt -> GenCtx ()
emit stmt = do
    (GenState stmts next) <- S.get
    S.put $ GenState (stmts ++ [stmt]) next

-- Get the name of the next temporary, and increment the value of the temp
-- counter
nextTemp :: GenCtx T.Text
nextTemp = S.state $
    \(GenState stmts next) ->
        (T.pack $ "__t" ++ show next, GenState stmts (next + 1))

runCtx :: GenCtx a -> Artifact [Stmt]
runCtx v = (S.runStateT v $ GenState [] 0) >>= \(_, GenState s _) -> return s

instance Show Atom where
    show (Name t) = T.unpack t
    show (InFloat f) = show f ++ "f"
    show (InInt i) = show i ++ "i"

instance Show IExpr where
    show (Add a1 a2) = show a1 ++ " + " ++ show a2
    show (Mul a1 a2) = show a1 ++ " * " ++ show a2
    show (Sub x y) = show x ++ " - " ++ show y
    show (Div x y) = show x ++ " / " ++ show y
    show (Lt x y) = show x ++ " < " ++ show y
    show (Gt x y) = show x ++ " > " ++ show y
    show (LtE x y) = show x ++ " <= " ++ show y
    show (GtE x y) = show x ++ " >= " ++ show y
    show (Negate x) = "~" ++ show x
    show (Ternary c t f) = show c ++ " ? " ++ show t ++ " : " ++ show f
    show (Atomic a) = show a
    show (Builtin n args)
        = T.unpack n ++ "(" ++ L.intercalate ", " (map show args) ++ ")"
    show (Index x i) = show x ++ "[" ++ show i ++ "]"
    show (Boolean b) = show b

instance Show Stmt where
    show (Store n t o) = "LET " ++ T.unpack n ++ ":" ++ show t ++ " = " ++ show o
    show (Return o) = "RET " ++ show o

instance Show Program where
    show (Program args ret stmts) =
        "args: " ++ show args ++ "\nreturns: " ++ show ret ++ "\n" ++
        L.intercalate "\n" (map show stmts)