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)

data CompResult = Result
  { -- last resolved type
    crType :: Maybe A.Type,
    crExpr :: A.Expr
  }

typeResult :: Maybe A.Type -> A.Expr -> CompResult
typeResult t e = Result t e

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 A.Nop

-- | @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
      -- FIXME: this break constant folding
      rargs <- traverse compileOne args
      let targs = map crType 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

typecheck :: Maybe A.Type -> Maybe A.Type -> Maybe String
typecheck expected found
  | expected == found = Nothing
  | otherwise = Just $ "type mismatch\n     found: " ++ showMaybet found ++ "\n  expected: " ++ showMaybet expected

-- 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 CompResult
compileOne x = do
  case x of
    (A.Module _ _) -> return $ typeResult Nothing x
    (A.Num _ _) -> return $ typeResult (Just $ A.Type "u8") x -- TODO: placeholder
    (A.Bool' _ _) -> return $ typeResult (Just $ A.Type "bool") x
    (A.BinOp A.Plus pos (A.Num a _) (A.Num b _)) ->
      -- TODO: overflow check
      return $ typeResult (Just $ A.Type "u8") (A.Num (a + b) pos)
    orig@(A.BinOp op pos a b) -> do
      ra <- compileOne a
      let (ta, ea) = (crType ra, crExpr ra)
      rb <- compileOne b
      let (tb, eb) = (crType rb, crExpr rb)
      case op of
        A.Assign -> case a of
          (A.Variable _ _) -> case typecheck ta tb of
            Nothing -> return $ typeResult ta (A.BinOp A.Assign pos ea eb)
            Just err -> addError $ Error TypeError err pos
          _ -> addError $ Error InvalidTarget "invalid assignment target" pos
        _ -> case typecheck ta tb of
          Nothing -> do
            let new = (A.BinOp op pos ea eb)
            if orig == new
              then return $ typeResult ta orig
              else compileOne new
          Just err -> addError $ Error TypeError err 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}
      r <- compileAll body
      st <- get
      -- store updated errors and the env with the function
      put st {stEnv = ev}
      rbody <- case r of
        Right (_, xs) -> pure $ xs
        Left _ -> pure $ body -- there was an error, so just keep the old body
      return $ typeResult (Just ftype) (A.Func ident params ret rbody priv anon pos)
      where
        ftype = A.toFuncType params ret
    (A.Call ident args pos) -> do
      r <- compileOne ident
      case crType 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 x
        _ -> 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}
      r <- compileOne val
      case typecheck (Just typ) (crType r) of
        Just err -> addError $ Error TypeError err pos
        Nothing -> return $ typeResult (Just typ) x
    (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 return (typeResult Nothing x)
                else addError $ Error TypeError ("invalid return value\n     found: ()\n  expected: " ++ showMaybet rtyp) pos
            Just v -> do
              r <- compileOne v
              case typecheck rtyp (crType r) of
                Just err -> addError $ Error TypeError err pos
                Nothing -> return $ typeResult rtyp (A.Return (Just (crExpr 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} -> return $ typeResult (Just t) x
        Nothing -> addError $ Error Undefined ("undefined \"" ++ ident ++ "\"") pos
    (A.Nop) -> return $ typeResult Nothing x

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

compileToAst :: [A.Expr] -> Either [Error] [A.Expr]
compileToAst ast = do
  (_, expr) <- evalState (compileAll ast) startState
  pure $ expr

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