content: module Lasca.EmitStatic where
  
  import LLVM.Module
  import LLVM.Context
  import LLVM.Analysis
  import LLVM.PassManager
  
  import qualified LLVM.AST as AST
  import qualified LLVM.AST.Global
  import qualified LLVM.AST.Type as T
  import qualified LLVM.AST.Instruction as I
  import qualified LLVM.AST.Constant as C
  import qualified LLVM.AST.Float as F
  import qualified LLVM.AST.IntegerPredicate as IP
  import qualified LLVM.AST.FloatingPointPredicate as FP
  import qualified LLVM.AST.FunctionAttribute as FA
  import qualified LLVM.AST.IntegerPredicate as IPred
  
  -- import qualified Data.Text as Text
  import qualified Data.ByteString as ByteString
  import qualified Data.Text.Encoding as Encoding
  import Text.Printf
  import qualified Data.ByteString.UTF8 as UTF8
  import Data.String
  import qualified Data.ByteString.Char8 as Char8
  import qualified Data.ByteString as BS
  import qualified Data.ByteString.Short as SBS
  
  import LLVM.ExecutionEngine ( withMCJIT, withModuleInEngine, getFunction )
  
  import qualified Data.Text
  import qualified Data.ByteString
  import qualified Data.Text.Encoding
  import Data.Digest.Murmur32
  import Data.Maybe
  import qualified Data.List as List
  import Data.Word
  import Data.Int
  import Control.Monad.State
  import Control.Monad.Except
  import Control.Applicative
  import qualified Control.Lens as Lens
  import Control.Lens.Operators
  import Data.Map.Strict (Map)
  import qualified Data.Map.Strict as Map
  import Data.Set (Set)
  import qualified Data.Set as Set
  import qualified Data.Sequence as Seq
  import qualified Debug.Trace as Debug
  import System.Exit
  import System.Directory
  import System.FilePath
  
  import Lasca.Codegen as Codegen
  import Lasca.Type
  import Lasca.EmitCommon
  import Lasca.Infer
  import qualified Lasca.Syntax as S
  import Lasca.Syntax (Ctx)
  import qualified Lasca.Options as Opts
  
  dataTypeHasField ctx typeName fieldName =
      typeName `Set.member` (S._dataDefsNames ctx) && fieldName `Map.member` (S._dataDefsFields ctx Map.! typeName)
  
  isDataType ctx tpe = case tpe of
      TypeIdent name | name `Set.member` (S._dataDefsNames ctx) -> True
      _ -> False
  
  isFuncType (TypeFunc _ _) = True
  isFuncType _ = False
  
  lascaPrimitiveTypes = Set.fromList [TypeInt, TypeFloat, TypeBool, TypeString, TypeAny, TypeUnit]
  anyTypeVar = TVar $ TV "a"
  
  cgen :: Ctx -> S.Expr -> Codegen AST.Operand
  cgen ctx (S.Let False meta a _ b c) = do
      i <- alloca $ llvmTypeOf b
      val <- cgen ctx b
      store i val
      assign a i
      cgen ctx c
  cgen ctx (S.Ident meta name) = do
      syms <- gets symtab
      modState <- gets moduleState
      let mapping = functions modState
      case lookup name syms of
          Just x ->
      --       Debug.trace ("Local " ++ show name)
              load x
          Nothing | name `Map.member` S._globalFunctions ctx -> boxClosure name mapping []
                  | name `Map.member` S._globalVals ctx -> load (globalOp ptrType (nameToSBS name))
                  | otherwise -> boxError (nameToText name)
  cgen ctx (S.Literal meta l) = do
  --  Debug.traceM $ "Generating literal " ++ show l ++ " on " ++ show (S.pos meta)
      boxLit l meta
  cgen ctx this@(S.Array meta exprs) = do
      vs <- sequence [cgen ctx e | e <- exprs]
      boxArray vs
  cgen ctx this@(S.Select meta tree expr) = cgenSelect ctx this
  cgen ctx this@(S.Apply meta (S.Ident _ "unary-") [expr]) = cgenApplyUnOp ctx this
  cgen ctx this@(S.Apply meta (S.Ident _ fn) [lhs, rhs]) | fn `Map.member` binops = cgenApplyBinOp ctx this
  -- TODO Either extend it to other types and generalize or make builtin functions inlineable
  cgen ctx this@(S.Apply meta (S.Ident _ (NS "Bits" fn)) [lhs, rhs]) | fn `elem` ["intAnd", "intOr", "intXor", "intShiftL", "intShiftR"] = do
      llhs <- cgen ctx lhs
      lrhs <- cgen ctx rhs
      a <- unboxInt llhs
      b <-  unboxInt lrhs
      case fn of
          (Name "intAnd") -> instrTyped intType (I.And a b []) >>= boxInt
          (Name "intOr")  -> instrTyped intType (I.Or a b []) >>= boxInt
          (Name "intXor") -> instrTyped intType (I.Xor a b []) >>= boxInt
          (Name "intShiftL") -> instrTyped intType (I.Shl False False a b []) >>= boxInt
          (Name "intShiftR") -> instrTyped intType (I.AShr False a b []) >>= boxInt
          _ -> error $ printf "Unsupported builtin operation %s" (show $ S.exprPosition this)
  
  cgen ctx (S.Apply meta expr args) = cgenApply ctx meta expr args
  cgen ctx (S.Closure _ funcName enclosedVars) = do
      modState <- gets moduleState
      let mapping = functions modState
      boxClosure funcName mapping enclosedVars
  cgen ctx m@S.Match{} =
      error $ printf "Match expressions should be already desugared! %s at: %s" (show m) (show $ S.exprPosition m)
  cgen ctx (S.If meta cond tr fl) = cgenIfStatic ctx meta cond tr fl
  
  cgen ctx e = error ("cgen shit " ++ show e)
  
  cgenIfStatic ctx meta cond tr fl = do
      let resultType = llvmTypeOf tr
      let test = do
              cond <- cgen ctx cond
              -- unbox Bool
              bool <- unboxBool cond
              instr (I.ICmp IP.EQ bool constTrue [])
      cgenIf resultType test (cgen ctx tr) (cgen ctx fl)
  
  cgenSelect ctx this@(S.Select meta tree expr) = do
      --    Debug.traceM $ printf "Selecting! %s" (show this)
      let (treeType, tpeName) = case S.typeOf tree of
                                    treeType@(TypeIdent tpeName) -> (treeType, tpeName)
                                    treeType@(TypeApply (TypeIdent tpeName) _) -> (treeType, tpeName)
                                    treeType -> error $ printf "Unsupported type for selection %s" (show treeType)
      let identType = S.typeOf expr
      let expectedReturnType = S.typeOf this
      --    Debug.traceM $ printf "Selecting %s: %s" (show treeType) (show identType)
      case expr of
         (S.Ident _ fieldName) | dataTypeHasField ctx tpeName fieldName -> do
             let pos = createPosition $ S.pos meta
             tree <- cgen ctx tree
             let (S.Ident _ fieldName) = expr
             let fieldsWithIndex = (S._dataDefsFields ctx) Map.! tpeName
      --            Debug.traceM $ printf "fieldsWithIndex %s" (show fieldsWithIndex)
             let (S.Arg n declaredFieldType, idx) = fromMaybe (error $ printf "No such field %s in %s" (show fieldName) (show tpeName)) (Map.lookup fieldName fieldsWithIndex)
             let len = length fieldsWithIndex
             dataStruct <- bitcast tree (T.ptr (dataValueStructType len))
             array <- getelementptr dataStruct [constIntOp 0, constInt32Op 2]
             valueAddr <- getelementptr array [constIntOp 0, constIntOp idx]
             load valueAddr
      --            traceM $ printf "AAAA %s: %s" (show array) (show value)
      --            resultValue <- castBoxedValue declaredFieldType value
      --            Debug.traceM $ printf "Selecting %s: %s" (show tree) (show resultValue)
      --            return $ constFloatOp 1234.5
      --            resolveBoxing declaredFieldType expectedReturnType resultValue
      --            return resultValue
          --    return value
         (S.Ident _ name) | isFuncType identType -> do
      --      traceM $ printf "Method call %s: %s" name (show identType)
             cgen ctx (S.Apply meta expr [tree])
         _ -> error $ printf "Unsupported select: %s at %s" (show this) (show $ S.pos meta)
  cgenSelect ctx e = error ("cgenSelect should only be called on Select, but called on" ++ show e)
  
  cgenApplyUnOp ctx this@(S.Apply meta op@(S.Ident _ "unary-") [expr]) = do
      lexpr' <- cgen ctx expr
      let (TypeFunc realExprType _) = S.typeOf op
      lexpr <- resolveBoxing anyTypeVar realExprType lexpr'
      let returnType = S.typeOf this
      let tpe = S.typeOf expr
      let minus = fromMaybe (error ("Only Byte | Int | Int16 | Int32 | Float supported but given " ++ show tpe)) (getArithOp 11 tpe)
      let llvmType = externalTypeMapping tpe
      let zero = case tpe of
              TypeByte -> C.Int 8 0
              TypeInt16 -> C.Int 16 0
              TypeInt32 -> C.Int 32 0
              TypeInt -> C.Int 64 0
              TypeFloat -> C.Float (F.Double 0.0)
              _ -> error $ printf "%s: Unexpected type in unaryOp - %s" (S.showPosition meta) (show tpe)
      r <- instrTyped llvmType $ (constOp zero) `minus` lexpr
      resolveBoxing returnType anyTypeVar r
  cgenApplyUnOp ctx e = error ("cgenApplyUnOp should only be called on Apply, but called on" ++ show e)
  
  getArithOp code tpe = case (code, tpe) of
      (10, _) | isIntegralType tpe -> Just $ \lhs rhs -> I.Add False False lhs rhs []
      (10, TypeFloat) -> Just $ \lhs rhs -> I.FAdd I.noFastMathFlags lhs rhs []
      (11, _) | isIntegralType tpe -> Just $ \lhs rhs -> I.Sub False False lhs rhs []
      (11, TypeFloat) -> Just $ \lhs rhs -> I.FSub I.noFastMathFlags lhs rhs []
      (12, _) | isIntegralType tpe -> Just $ \lhs rhs -> I.Mul False False lhs rhs []
      (12, TypeFloat) -> Just $ \lhs rhs -> I.FMul I.noFastMathFlags lhs rhs []
      (13, _) | isIntegralType tpe -> Just $ \lhs rhs -> I.SDiv True lhs rhs []
      (13, TypeFloat) -> Just $ \lhs rhs -> I.FDiv I.noFastMathFlags lhs rhs []
      _ -> Nothing
  
  getCmpOp code tpe = case (code, tpe) of
      (42, _) | isIntegralType tpe || tpe == TypeBool -> Just $ \lhs rhs -> I.ICmp IPred.EQ lhs rhs []
      (42, TypeFloat) -> Just $ \lhs rhs -> I.FCmp FP.OEQ lhs rhs []
      (43, _) | isIntegralType tpe || tpe == TypeBool -> Just $ \lhs rhs -> I.ICmp IPred.NE lhs rhs []
      (43, TypeFloat) -> Just $ \lhs rhs -> I.FCmp FP.ONE lhs rhs []
      (44, _) | isIntegralType tpe || tpe == TypeBool -> Just $ \lhs rhs -> I.ICmp IPred.SLT lhs rhs []
      (44, TypeFloat) -> Just $ \lhs rhs -> I.FCmp FP.OLT lhs rhs []
      (45, _) | isIntegralType tpe || tpe == TypeBool -> Just $ \lhs rhs -> I.ICmp IPred.SLE lhs rhs []
      (45, TypeFloat) -> Just $ \lhs rhs -> I.FCmp FP.OLE lhs rhs []
      (46, _) | isIntegralType tpe || tpe == TypeBool -> Just $ \lhs rhs -> I.ICmp IPred.SGE lhs rhs []
      (46, TypeFloat) -> Just $ \lhs rhs -> I.FCmp FP.OGE lhs rhs []
      (47, _) | isIntegralType tpe || tpe == TypeBool -> Just $ \lhs rhs -> I.ICmp IPred.SGT lhs rhs []
      (47, TypeFloat) -> Just $ \lhs rhs -> I.FCmp FP.OGT lhs rhs []
      _ -> Nothing
  
  cgenApplyBinOp ctx this@(S.Apply meta op@(S.Ident _ fn) [lhs, rhs]) = do
      llhs' <- cgen ctx lhs
      lrhs' <- cgen ctx rhs
      let lhsType = S.typeOf lhs
      let rhsType = S.typeOf rhs
      let returnType = S.typeOf this
  --    Debug.traceM $ printf "Doing binop %s with type %s" (show this) (show $ S.typeOf op)
      let (realLhsType, realRhsType) = case S.typeOf op of
                  TypeFunc realLhsType (TypeFunc realRhsType _) -> (realLhsType, realRhsType)
                  _ -> error ("cgenApplyBinOp: Should not happen: " ++ show this ++ show (S.typeOf op))
      llhs <- resolveBoxing anyTypeVar realLhsType llhs'
      lrhs <- resolveBoxing anyTypeVar realRhsType lrhs'
      let code = fromMaybe (error ("Couldn't find binop " ++ show fn)) (Map.lookup fn binops)
  --    Debug.traceM $ printf "%s: %s <==> %s: %s, code %s" (show lhsType) (show realLhsType) (show rhsType) (show realRhsType) (show code)
      let llvmType = externalTypeMapping realLhsType
      if isPrimitiveType realLhsType
      then do
          res <- case code of
              _ | code >= 10 && code <= 13 -> do
                  let op = fromMaybe (error $ printf "cgenApplyBinOp not defined operation code %d for type %s" code (show realLhsType)) (getArithOp code realLhsType)
                  instrTyped llvmType (llhs `op` lrhs)
              _ | code >= 42 && code <= 47 -> do
                  let op = fromMaybe (error $ printf "cgenApplyBinOp not defined operation code %d for type %s" code (show realLhsType)) (getCmpOp code realLhsType)
                  r <- instrTyped llvmType (llhs `op` lrhs)
                  instrTyped boolType $ I.ZExt r boolType []
              c  -> error $ printf "%s: Unsupported binary operation %s, code %s, type %s" (show $ S.exprPosition this) (S.printExprWithType this) (show c) (show realLhsType)
          resolveBoxing returnType anyTypeVar res
      else callBuiltin "runtimeBinOp" [constIntOp code, llhs', lrhs']
  cgenApplyBinOp ctx e = error ("cgenApplyBinOp should only be called on Apply, but called on" ++ show e)
  
  cgenApply ctx meta expr args = do
      syms <- gets symtab
      let symMap = Map.fromList syms
      let isGlobal fn = (fn `Map.member` S._globalFunctions ctx) && not (fn `Map.member` symMap)
      case expr of
           -- FIXME Here are BUGZZZZ!!!! :)
          this@(S.Ident meta (NS "Array" "getIndex")) -> do
              let [arrayExpr, indexExpr] = args
              array <- cgen ctx arrayExpr -- should be a pointer to either boxed or unboxed array
              boxedIdx <- cgen ctx indexExpr
              idx <- unboxInt boxedIdx
  --            callFn (funcType ptrType [ptrType, intType]) "arrayGetIndex" [array, idx]
              cgenArrayApply array idx
  
          S.Ident _ fn | isGlobal fn -> do
  --            Debug.traceM $ printf "Calling %s" fn
              let f = S._globalFunctions ctx Map.! fn
              largs <- forM args $ \arg -> cgen ctx arg
              call (funcLLvmType f) (nameToSBS fn) largs
          expr -> do
              -- closures
              modState <- gets moduleState
              e <- cgen ctx expr
              largs <- mapM (cgen ctx) args
              let argc = constIntOp (length largs)
              sargsPtr <- allocaSize ptrType argc
              let asdf (idx, arg) = do
                    p <- getelementptr sargsPtr [idx]
                    store p arg
              sargs <- bitcast sargsPtr ptrType -- runtimeApply accepts i8*, so need to bitcast. Remove when possible
              -- cdecl calling convension, arguments passed right to left
              sequence_ [asdf (constIntOp i, a) | (i, a) <- zip [0..] largs]
  
              let pos = createPosition $ S.pos meta
              callBuiltin "runtimeApply" [e, argc, sargs, constOp pos]
  
  cgenArrayApply array idx = do
      arrayStructPtr <- bitcast array (T.ptr $ arrayStructType ptrType) -- &Array(type, len, &data[])
      -- TODO check idx is in bounds, eliminatable
      ptr <- getelementptr arrayStructPtr [constIntOp 0, constInt32Op 2, idx]
      load ptr
  
  -------------------------------------------------------------------------------
  -- Compilation
  -------------------------------------------------------------------------------
  
  funcPtrFromClosure closure = do
      modState <- gets moduleState
      let mapping = functions modState
      let len = Map.size mapping
      closureTyped <- bitcast closure (T.ptr closureStructType)
      idxPtr <- getelementptr closureTyped [constIntOp 0, constInt32Op 1]
      idx <- instrTyped intType $ I.Load False idxPtr Nothing 0 []
  --    callFn "putInt" [idx]
      let fst = functionsStructType (fromIntegral len)
      let fnsAddr = (globalOp fst (nameToSBS "Functions"))
      fns <- instrTyped (T.ptr fst) $ I.Load False fnsAddr Nothing 0 []
  --    sizeAddr <- getelementptr fnsAddr [constIntOp 0, constIntOp 0]
  --    size <- instrTyped intType $ I.Load False sizeAddr Nothing 0 []
  --     Functions[idx].funcPtr
      fnPtr <- getelementptr fnsAddr [constIntOp 0, constInt32Op 1, idx, constInt32Op 1]
      load fnPtr
  
  resolveBoxing declaredType instantiatedType expr =
      case (declaredType, instantiatedType) of
          _ | declaredType == instantiatedType -> return expr
          (TypeByte, TVar _) -> boxByte expr
          (TypeBool, TVar _) -> boxBool expr
          (TypeInt, TVar _) -> boxInt expr
          (TypeInt16, TVar _) -> boxInt16 expr
          (TypeInt32, TVar _) -> boxInt32 expr
          (TypeFloat, TVar _) -> boxFloat64 expr
          (TVar _, TypeByte) -> unboxByte expr
          (TVar _, TypeBool) -> unboxBool expr
          (TVar _, TypeInt) -> unboxInt expr
          (TVar _, TypeInt16) -> unboxInt16 expr
          (TVar _, TypeInt32) -> unboxInt32 expr
          (TVar _, TypeFloat) -> unboxFloat64 expr
          (TVar _, TVar _) -> return expr
          (l, r) -> do
  --            Debug.traceM $ printf "resolveBoxing crap %s %s" (show l) (show r)
              return expr
  {-# INLINE resolveBoxing #-}
  
  
  