module Compiler where

import Data.List
import Syntax.AST

import Control.Monad.State
import Control.Monad.Writer

import Prelude hiding (mod)
import Debug.Trace

debug = flip trace

compilePath :: Path -> String
compilePath (Path xs) = intercalate "." $ filter (not . null) xs

data CompilerState = CompilerState { ident :: Int, mod :: String, exps :: [(String, String)] }

type Compiler = StateT CompilerState (Writer String) ()

block :: Compiler -> Compiler
block x = do
  modify $ \(CompilerState x y z) -> CompilerState (x + 1) y z 
  x
  modify $ \(CompilerState x y z) -> CompilerState (x - 1) y z

incindent :: Compiler
incindent
  = modify $ \(CompilerState x y z) -> CompilerState (x + 1) y z 

decindent :: Compiler
decindent
  = modify $ \(CompilerState x y z) -> CompilerState (x - 1) y z 

parens :: Compiler -> Compiler
parens x = do
  tell "("
  x
  tell ")"

newline = tell "\n"

indent :: Compiler -> Compiler
indent next = do
  x <- gets ident
  tell $ concat $ replicate x "  "
  next 

compileAssignment :: [Name] -> [Expression] -> Compiler
compileAssignment x y
  = tell $ intercalate ", " x ++ " = " ++ intercalate ", " (map compileLoneExpr y)


compileStatement :: Statement -> Compiler
compileStatement (SExpression ex) = compileExpression ex
compileStatement (SReturn ex) = tell "return " >> compileExpression ex
compileStatement (SAssign s e) = tell "local " >> compileAssignment s e 
compileStatement (SReassign s e) = compileAssignment s e

compileStatement (SFunction (n,_) params []) = do
  tell $ "function " ++ n
  parens $ tell $ intercalate ", " $ (\(x, _, _) -> x) `map` params
  tell " end"
  newline

compileStatement (SFunction (n,_) params body) = do
  tell $ "function " ++ n
  parens $ tell $ intercalate ", " $ (\(x, _, _) -> x) `map` params
  block $ do mapM_ ((>> newline) . indent . makeDefault) params
             compileBody body
  newline
  tell "end"

compileStatement (SOpen e@(Path mod)) = do
  tell $ "local " ++ last mod ++ " = require('" ++ compilePath e ++ "')"
  newline
compileStatement (SFFI (Path mod) nam _) = do
  tell "local "
  tell nam
  tell " = "
  tell $ intercalate "." mod
  newline

compileStatement (SStruct nam fields) = do
  tell "function "
  tell nam
  parens $ tell $ intercalate ", " $ map fst fields
  newline
  block $ do tell $ "return { tag = " ++ show nam ++ ", "
             indent $ tell $ intercalate ", " $ map gen fields
             tell "} "
  tell "end"
  newline
    where gen = (\x -> x ++ " = " ++ x) . fst

compileStatement (STypeAlias _ _) = return ()
compileStatement (SModule mn) = do
  tell $ "local " ++ mn ++ " = {} "
  modify $ \x -> x { mod = intercalate "." $ filter (not . null) [mod x, mn] } 
  newline

compileStatement (SExport fun@(SFunction (nam, _) _ _)) = do
  mnam <- gets mod
  modify $ \x -> x { exps = (mnam, nam):exps x }
  compileStatement fun


iife :: Compiler -> Compiler
iife x = do
  tell "(function()\n"
  block x
  newline
  indent $ tell "end)()"

makeDefault :: (String, Maybe Type, Maybe Expression) -> Compiler
makeDefault (x, _, Just y) = tell ("local " ++ x ++ " = " ++ x ++ " or ") >> compileExpression y
makeDefault (x, _, Nothing) = return () 

compileBody x = do block $ mapM_ ((>> newline) . indent . compileStatement) $ init x
                   let xs = last x
                   indent $ compileStatement $ case xs of
                                        (SReturn e)   -> SReturn e
                                        (SExpression e) -> SReturn e
                                        _              -> xs
                   tell "; "

compileExpression :: Expression -> Compiler
compileExpression (ELiteral s) = tell $ compileLiteral s
compileExpression (EIf cs e) = iife $ do
  compileSingle False $ head cs
  mapM_ (compileSingle True) $ tail cs
  newline
  case e of
    ls@(x:xs) -> indent $ tell "else" >> newline >> block (compileBody ls) >> newline
    []        -> return ()
  indent $ tell "end"
    where compileSingle x (t, st) = indent $ do tell $ if x then " elseif " else "if "
                                                compileExpression t
                                                tell " then "
                                                newline
                                                block $ compileBody st

compileExpression (EIndex e (ETuple _)) = error "Can not index an expression with a tuple."
compileExpression (EIndex (EVar e) (EVar x))
  = tell $ concat [e, ".", x]

compileExpression (EIndex e x) = do
  parens $ compileExpression e
  tell "["
  compileExpression x
  tell "]"

compileExpression (ELambda v []) =  do
  tell "function"
  parens $ tell $ intercalate ", " $ map (\(x, _, _) -> x) v
  indent $ tell " end"
compileExpression (ELambda v e) =  do
  tell "function"
  parens $ tell $ intercalate ", " $ map (\(x, _, _) -> x) v
  newline
  block $ do forM_ v $ \s -> do indent $ makeDefault s
                                newline
             compileBody e
  newline
  indent $ tell "end"

compileExpression (EFor n it bd)
  = iife $ do indent $ tell "local __r, __i = {}, 0\n"
              x <- get
              indent $ tell "for "
              tell $ intercalate ", " n
              tell " in "
              put $ CompilerState 0 (mod x) (exps x)
              compileExpression it
              put x
              tell " do\n"
              block $ do
                forM_ (init bd) $ \x -> indent $ compileStatement x >> tell "\n"
                indent $ tell "__i = __i + 1\n"
                let ls = last bd in
                    case ls of 
                      SExpression e  -> indent $ tell "__r[__i] = " >> compileExpression e >> tell "\n"
                      SReturn e      -> indent $ tell "__r[__i] = " >> compileExpression e >> tell "\n"
                      _              -> indent $ compileStatement ls >> newline
              indent $ tell "end" >> newline
              indent $ tell "return __r"

compileExpression (EWhile cn bd)
  = iife $ do 
      indent $ tell "local __r, __i = {}, 0\n"
      x <- get
      indent $ tell "while "
      put $ CompilerState 0 (mod x) (exps x)
      compileExpression cn
      tell " do\n"
      put x
      block $ do mapM_ ((>> newline) . indent . compileStatement) $ init bd
                 indent $ tell "__i = __i + 1\n"
                 let ls = last bd in
                     case ls of 
                       SExpression e  -> indent $ tell "__r[__i] = " >> compileExpression e >> newline
                       SReturn e      -> indent $ tell "__r[__i] = " >> compileExpression e >> newline
                       _              -> indent $ compileStatement ls >> newline
      indent $ tell "end\n"
      indent $ tell "return __r"

compileExpression (EList es) = do
  tell "{ "
  forM_ es $ \s -> compileExpression s >> tell ", "
  tell " }"

compileExpression (EMatch bd e) = do
  tell "(function (x)\n"
  block $ indent $ compileStatement $ buildExprs bd
  tell "\nend)("
  compileExpression e
  tell ")"
    where buildExprs x = SReturn $ EIf (map buildSingle x) [SReturn $ ELiteral LUnit]
          buildSingle (x, y) = (EBinOp Eq (EVar "x") x, [SReturn y])


compileExpression (ETuple es) = tell $ intercalate ", " $ map compileLoneExpr es

compileExpression (ECall (EVar x) xs) = do
  tell x
  let argsl = intercalate ", " $ map compileLoneExpr xs in
      parens $ tell argsl

compileExpression (ECall e xs) = do
  parens $ compileExpression e
  let argsl = intercalate ", " (map compileLoneExpr xs) in
    parens $ tell argsl

compileExpression (EVar name) = tell name
compileExpression (EPath path) = tell $ compilePath path
compileExpression (EBinOp op lh rh) = do
  compileExpression lh
  tell $ case op of
            Add -> " + " 
            Sub -> " - "
            Div -> " / "
            Mul -> " * "
            And -> " and "
            Or -> " or "
            Eq -> " == "
            Shl -> " << "
            Shr -> " >> "
            Lt -> " < "
            Le -> " <= "
            Gt -> " > "
            Ge -> " >= "
            Concat -> " .. "
            Pow -> " ^ "
            BitOr -> " & "
            BitAnd -> " | "
            Neq -> " ~= "
            Mod -> " % "
  compileExpression rh

compileExpression (EUnOp op e) = do
  tell $ case op of
           Len -> " #"
           BitNot -> " ~"
           Not -> " not "
  compileExpression e

compileExpression (EStructConstructor p es) = 
  compileExpression $ ECall (EPath p) (map snd es)

compileExpression (EListComprehend lhs rhses)
  = iife $ compileListComprehension lhs rhses

compileExpression (EListRange s t (Just inc))
  = iife $ do indent $ tell "local _ = {}\n"
              indent $ tell "local st, en, sp = "
              compileExpression s
              tell ", "
              compileExpression t
              tell ", "
              compileExpression inc
              newline
              indent $ tell "for i = st, en, (sp - st) do\n"
              block $ indent $ tell "_[#_ + 1] = i\n"
              indent $ tell "end"
              indent $ tell "return _"

compileExpression (EListRange s t Nothing)
  = iife $ do indent $ tell "local _ = {}\n"
              indent $ tell "local st, en = "
              compileExpression s
              tell ", "
              compileExpression t
              newline
              indent $ tell "for i = st, en do\n"
              block $ indent $ tell "_[#_ + 1] = i\n"
              indent $ tell "end\n"
              indent $ tell "return _"

compileLiteral LUnit = "nil"
compileLiteral (LInteger i) = show i
compileLiteral (LDouble i) = show i
compileLiteral (LString i) = show i
compileLiteral (LBoolean True) = "true"
compileLiteral (LBoolean False) = "false"

compileStatements x = do
  mapM_ compileStatement x
  x <- gets mod
  exp <- gets exps
  forM_ exp $ \(m, x) -> tell $ intercalate "." [m, x] ++ " = " ++ x ++ "\n"

  tell "if main then arg = {...}; main(#arg, arg); end;"
  tell $ if not (null x) then "return " ++ x else ""

emptyCS = CompilerState 0 "" []


compileTop :: [Statement] -> String
compileTop xs = let ((_, _), x) = runWriter (compileStatements xs `runStateT` emptyCS) in x

compileLoneExpr :: Expression -> String
compileLoneExpr x = let ((_, _), x') = runWriter (compileExpression x `runStateT` emptyCS) in x'

compileListComprehension :: Expression -> [Statement] -> Compiler
compileListComprehension ex st
  = do indent $ tell "local ret = {}\n"
       bounds' <- forM st $ \s -> case s of
           e@SAssign{} -> return $ Just e
           _           -> return Nothing
       tests' <- forM st $ \s -> case s of
           e@SExpression{} -> return $ Just e
           _               -> return Nothing
       let bounds = map (\(Just e) -> e) $ filter isJust bounds'
       let tests  = map (\(Just e) -> e) $ filter isJust tests'
       forM_ bounds
           $ \(SAssign [n] [e]) -> do
             indent $ tell $ concat [ "for _, ", n
                                    , " in pairs("]
             block $ compileExpression e 
             tell $ ") do\n"
             incindent
       forM_ tests
           $ \(SExpression e) ->
             do indent $ tell "if "
                compileExpression e
                tell " then\n"
                incindent

       indent $ tell "ret[#ret + 1] = "
       compileExpression ex
       newline
       forM_ bounds $ \_ -> do {decindent; indent $ tell "end\n"}
       forM_ tests $ \_ -> do {decindent; indent $ tell "end\n"}
       indent $ tell "return ret"


isJust :: Maybe a -> Bool
isJust Just{} = True
isJust      _ = False
