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

import Control.Arrow (ArrowChoice (left))
import Control.Monad.State
import Data.Foldable (traverse_)
import Data.Maybe (isNothing, mapMaybe)
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 =
  foldl
    ( \(r, l) x -> case fn r x of
        Left e -> (r, e : l)
        Right x -> (x, l)
    )

-- | @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
  modify $ \st -> 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"
  | null params = 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 -> pure Nothing
          -- compare types
          | all (uncurry (==)) $ zip t params ->
              pure 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 && maybe False definedType r

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

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 _ pos (A.Bool' _ _) (A.Bool' _ _)) ->
      -- for now this is true
      addError $ Error InvalidOperation "invalid operation for type bool" pos
    (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)
      let fev = 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
      -- store updated errors and the env with the function
      modify $ \st -> 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 local pos) -> do
      st <- get
      (ev, errs) <- return $ foldlEither addSymUniq (stEnv st, stErr st) [Sym ident typ priv local 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 -> Either Error A.Expr
foldConstant x =
  case x of
    -- FIXME: overflow?
    (A.BinOp A.Plus pos (A.Num a _) (A.Num b _)) -> Right $ A.Num (a + b) pos
    (A.BinOp A.Minus pos (A.Num a _) (A.Num b _)) -> Right $ A.Num (a - b) pos
    (A.BinOp A.Mul pos (A.Num a _) (A.Num b _)) -> Right $ A.Num (a * b) pos
    (A.BinOp A.Mul pos _ (A.Num 0 _)) -> Right $ A.Num 0 pos
    (A.BinOp A.Mul pos (A.Num 0 _) _) -> Right $ A.Num 0 pos
    (A.BinOp A.Div pos _ (A.Num 0 _)) -> Left $ Error InvalidOperation "division by zero" pos
    (A.BinOp A.Div pos (A.Num a _) (A.Num b _)) -> Right $ A.Num (a `div` b) pos
    (A.BinOp op pos a b) -> do
      fa <- foldConstant a
      fb <- foldConstant b
      let newOp = A.BinOp op pos fa fb
      if newOp /= x then foldConstant newOp else Right newOp
    (A.Func ident params ret body priv anon pos) -> do
      fbody <- traverse foldConstant body
      Right $ A.Func ident params ret fbody priv anon pos
    (A.Call ident args pos) -> do
      fid <- foldConstant ident
      fargs <- traverse foldConstant args
      Right $ A.Call fid fargs pos
    (A.Var ident typ val priv local pos) -> do
      fv <- foldConstant val
      Right $ A.Var ident typ fv priv local pos
    (A.Return value pos) -> do
      fv <- traverse foldConstant value
      Right $ A.Return fv pos
    _ -> Right 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 = do
  _ <- evalState (compileAll ast) startState
  left (: []) $ traverse foldConstant ast

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