module Micro.Compiler where

import Control.Monad.State
import Data.Maybe (catMaybes, fromMaybe)
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],
    stAst :: [A.Expr]
  }
  deriving (Show)

data CompResult = CompResult
  { -- last resolved type
    crLast :: Maybe A.Type,
    -- only in last call
    crExit :: Either [Error] String
  }

typeResult :: Maybe A.Type -> CompResult
typeResult t = CompResult {crLast = t, crExit = Left []}

errorResult :: [Error] -> CompResult
errorResult err = CompResult {crLast = Nothing, crExit = Left err}

successResult :: String -> CompResult
successResult out = CompResult {crLast = Nothing, crExit = Right out}

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 CompResult
addError e = do
  st <- get
  put st {stErr = e : stErr st}
  pure $ typeResult 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
      rargs <- traverse compile args
      let targs = map crLast rargs
      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

showMaybet :: Maybe A.Type -> String
showMaybet Nothing = "()"
showMaybet (Just t) = show t

-- | @typecheckVal value typ@ resolves @value@ and compares it to @typ@ type,
-- returning a string describing an error or Nothing in case of type match.
typecheckVal :: A.Expr -> Maybe A.Type -> State CompState (Maybe String)
typecheckVal value typ = do
  r <- compile value
  case crLast r of
    rt
      | rt == typ -> pure Nothing
      | otherwise -> return $ Just $ "type mismatch\n     found: " ++ showMaybet rt ++ "\n  expected: " ++ showMaybet typ

-- | @typecheckReturn value fret@ resolves @value@ and compares it with @fret@,
-- returning a string decribing an error or Nothing in case of a type match.
typecheckReturn :: Maybe A.Expr -> Maybe A.Type -> State CompState (Maybe String)
typecheckReturn Nothing Nothing = return $ Nothing
typecheckReturn Nothing fret = return $ Just $ "invalid return value\n     found: ()\n  expected: " ++ showMaybet fret
typecheckReturn (Just value) fret = typecheckVal value fret

-- | @typecheckBinOp a b pos@ resolves @a@ (left) and compares it to the type
-- of @b@ via typecheckVal.
typecheckBinOp :: A.Expr -> A.Expr -> SourcePos -> State CompState CompResult
typecheckBinOp a b pos = do
  l <- compile a
  let tl = crLast l
  tr <- typecheckVal b $ tl
  case tr of
    Just err -> addError $ Error TypeError err pos
    Nothing -> return $ typeResult tl

-- 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 []

compile :: A.Expr -> State CompState CompResult
compile x = do
  case x of
    (A.Module _ _) -> return $ typeResult Nothing
    (A.Num _ _) -> return $ typeResult $ Just $ A.Type "u8" -- TODO: placeholder
    (A.Bool' _ _) -> return $ typeResult $ Just $ A.Type "bool"
    (A.BinOp A.Assign pos a@(A.Variable _ _) b) ->
      typecheckBinOp a b pos
    (A.BinOp A.Assign pos _ _) ->
      addError $ Error InvalidTarget "invalid assignment target" pos
    (A.BinOp _ pos a b) ->
      -- TODO: types and invalid operators
      typecheckBinOp a b 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, stErr = errs}
      return $ typeResult $ Just ftype
      where
        ftype = A.toFuncType params ret
    (A.Call ident args pos) -> do
      r <- compile ident
      case crLast r of
        Just (A.FuncType params rtyp) -> do
          r <- typecheckCall args params
          case r of
            Just err -> addError $ Error TypeError err pos
            Nothing -> return $ typeResult 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}
      vt <- typecheckVal val $ Just typ
      case vt of
        Just err -> addError $ Error TypeError err pos
        Nothing -> return $ typeResult $ Just typ
    (A.Return value pos) -> do
      st <- get
      case getSyml (stEnv st) "$fn$" of
        Just Sym {symType = A.FuncType _ rtyp} -> do
          r <- typecheckReturn value rtyp
          case r of
            Just err -> addError $ Error TypeError err pos
            Nothing -> return $ typeResult rtyp
        _ -> 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} -> return $ typeResult $ Just t
        Nothing -> addError $ Error Undefined ("undefined \"" ++ ident ++ "\"") pos

compileAll :: [A.Expr] -> State CompState CompResult
compileAll ast = do
  _ <- traverse compile ast
  st <- get
  let (Env sym _) = stEnv st
  let errs = stErr st
  case errs of
    [] -> pure $ successResult $ generate version sym ast
    _ -> return $ errorResult errs