module Micro.Compiler
  ( version,
    compile,
    compileToAst,
  )
where

import Control.Monad.State
import Data.Maybe (catMaybes, fromMaybe, isNothing)
import Micro.Asm.Sdcc (generate)
import qualified Micro.Ast as A
import Micro.Env
import Micro.Error
import Text.Parsec (SourcePos)

version :: String
version = "0.1.0"

data CompState = CompState
  { stEnv :: Env,
    stErr :: [Error]
  }
  deriving (Show)

startState :: CompState
startState = CompState emptyEnv []

-- | @foldlEither fn init xs@ folds left @xs@ applying a function @fn@ that
-- returns either, returning accumulated right and the collected lefts as a
-- list.
foldlEither :: (accr -> b -> Either accl accr) -> (accr, [accl]) -> [b] -> (accr, [accl])
foldlEither fn init xs =
  foldl
    ( \(r, l) x -> case fn r x of
        Left e -> (r, e : l)
        Right x -> (x, l)
    )
    init
    xs

-- | @addError error@ adds @error@ to the state and returns no type to allow
-- the compilation to continue.
addError :: Error -> State CompState (Maybe A.Type)
addError e = do
  st <- get
  put st {stErr = e : stErr st}
  pure $ Nothing

-- | @typecheckCall args params@ resolves @args@ and compares it with @params@,
-- returning a string describing an error or Nothing in case of type match.
typecheckCall :: [A.Expr] -> [A.Type] -> State CompState (Maybe String)
typecheckCall args params
  | length args /= length params = return $ Just "invalid number of arguments in function call"
  | length params == 0 = pure Nothing
  | otherwise = do
      -- resolve all args types
      targs <- traverse compileOne args
      case sequence targs of
        Just t
          -- after resolution, we could have less "rights"
          | length t /= length params -> return $ Nothing
          -- compare types
          | all (\(a, b) -> a == b) $ zip t params ->
              return $ Nothing -- all good!
          | otherwise -> return $ Just ("type mismatch in function call\n     found: " ++ A.showList t ++ "\n  expected: " ++ A.showList params)
        Nothing ->
          -- there was an error in on argument
          pure Nothing

-- | @showMaybeType t@ is a helper to show nicely a @Maybe A.Type@.
showMaybeType :: Maybe A.Type -> String
showMaybeType Nothing = "()"
showMaybeType (Just t) = show t

-- | @typecheck expected found pos@ compares expected and found, returning either expected of adding an type error to the state.
typecheck :: Maybe A.Type -> Maybe A.Type -> SourcePos -> State CompState (Maybe A.Type)
typecheck expected found pos
  | expected == found = pure expected
  | otherwise = addError $ Error TypeError ("type mismatch\n     found: " ++ showMaybeType found ++ "\n  expected: " ++ showMaybeType expected) pos

-- built-in types
types :: [String]
types = ["bool", "u8", "s8", "u16", "s16"]

definedType :: A.Type -> Bool
definedType (A.Type t) = t `elem` types
definedType (A.FuncType ts r) =
  all definedType ts && fromMaybe False (fmap definedType r)

verifyFuncType :: String -> [A.FuncParam] -> Maybe A.Type -> SourcePos -> [Error]
verifyFuncType ident params ret pos = do
  ( catMaybes $
      map
        ( \(_, t, _, pos) ->
            if not (definedType t)
              then Just $ Error UndefinedType ("undefined type in declaration of \"" ++ ident ++ "\"") pos
              else Nothing
        )
        params
    )
    ++ if not (fromMaybe True (fmap definedType ret))
      then [Error UndefinedType ("undefined return type in \"" ++ ident ++ "\"") pos]
      else []

compileOne :: A.Expr -> State CompState (Maybe A.Type)
compileOne x = do
  case x of
    (A.Module _ _) -> pure Nothing
    (A.Num _ _) -> pure $ Just $ A.Type "u8" -- TODO: placeholder
    (A.Bool' _ _) -> pure $ Just $ A.Type "bool"
    (A.BinOp op pos a b) -> do
      ta <- compileOne a
      tb <- compileOne b
      case op of
        A.Assign -> case a of
          (A.Variable _ _) -> typecheck ta tb pos
          _ -> addError $ Error InvalidTarget "invalid assignment target" pos
        _ -> typecheck ta tb pos
    (A.Func ident params ret body priv anon pos) -> do
      -- current env
      st <- get
      -- check for undefined types
      (ev, errs) <- return (stEnv st, (verifyFuncType ident params ret pos) ++ stErr st)
      -- updated with the function
      (ev, errs) <-
        return $ case addSymUniq ev (Sym ident ftype priv False pos) of
          Left err -> (ev, err : errs)
          Right ev -> (ev, errs)
      -- lambdas can only access local variables (closures aren't supported)
      fev <- return $ if anon then emptyEnv else ev
      -- with parameters
      (nev, errs) <- return $ foldlEither addSymUniq (addEnv fev, errs) $ map (toSym True) params
      -- helper for return
      nev <- return $ addSym nev $ Sym "$fn$" ftype True True pos
      put st {stEnv = nev, stErr = errs}
      _ <- compileAll body
      st <- get
      -- store updated errors and the env with the function
      put st {stEnv = ev}
      pure Nothing
      where
        ftype = A.toFuncType params ret
    (A.Call ident args pos) -> do
      r <- compileOne ident
      case r of
        Just (A.FuncType params rtyp) -> do
          e <- typecheckCall args params
          case e of
            Just err -> addError $ Error TypeError err pos
            Nothing -> pure rtyp
        _ -> addError $ Error NonCallable "non callable value in function call" pos
    (A.Var ident typ val priv pos) -> do
      st <- get
      (ev, errs) <- return $ foldlEither addSymUniq (stEnv st, stErr st) [Sym ident typ priv True pos]
      errs <-
        return $
          if not (definedType typ)
            then Error UndefinedType ("undefined type in declaration \"" ++ ident ++ "\"") pos : errs
            else errs
      put st {stEnv = ev, stErr = errs}
      t <- compileOne val
      typecheck (Just typ) t pos
    (A.Return value pos) -> do
      st <- get
      case getSyml (stEnv st) "$fn$" of
        Just Sym {symType = A.FuncType _ rtyp} -> do
          case value of
            Nothing ->
              if isNothing rtyp
                then pure $ Nothing
                else addError $ Error TypeError ("invalid return value\n     found: ()\n  expected: " ++ showMaybeType rtyp) pos
            Just v -> do
              r <- compileOne v
              typecheck rtyp r pos
        _ -> addError $ Error UnexpectedReturn "return without function call" pos
    (A.Variable ident pos) -> do
      st <- get
      case getSym (stEnv st) ident of
        Just Sym {symType = t} -> pure $ Just t
        Nothing -> addError $ Error Undefined ("undefined \"" ++ ident ++ "\"") pos

foldConstant :: A.Expr -> A.Expr
foldConstant x =
  case x of
    (A.BinOp A.Plus pos (A.Num a _) (A.Num b _)) -> A.Num (a + b) pos
    (A.BinOp op pos a b) -> do
      let newOp = A.BinOp op pos (foldConstant a) (foldConstant b)
      if newOp /= x then foldConstant newOp else newOp
    (A.Func ident params ret body priv anon pos) ->
      A.Func ident params ret (map foldConstant body) priv anon pos
    (A.Call ident args pos) ->
      A.Call (foldConstant ident) (map foldConstant args) pos
    (A.Var ident typ val priv pos) -> A.Var ident typ (foldConstant val) priv pos
    (A.Return value pos) -> A.Return (fmap foldConstant value) pos
    _ -> x

compileAll :: [A.Expr] -> State CompState (Either [Error] SymMap)
compileAll ast = do
  _ <- traverse compileOne ast
  st <- get
  let (Env sym _) = stEnv st
  let errs = stErr st
  case errs of
    [] -> pure $ Right sym
    _ -> return $ Left errs

compileToAst :: [A.Expr] -> Either [Error] [A.Expr]
compileToAst ast = fmap (\_ -> map foldConstant ast) (evalState (compileAll ast) startState)

compile :: [A.Expr] -> Either [Error] String
compile ast = do
  sym <- evalState (compileAll ast) startState
  return $ generate version sym $ map foldConstant ast