module Compiler where

import qualified Ast as A
import Control.Monad.State
import Data.Either (rights)
import Env
import Error
import System.Environment (getEnv, getEnvironment)
import Text.Parsec (ParseError, SourcePos)
import Text.Read (Lexeme (String))

type CompState = (Env, [Error])

type CompResult = Either [Error] (Maybe A.Type)

startState :: CompState
startState = (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
  (ev, errs) <- get
  put (ev, e : errs)
  return $ Right 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 = return $ Nothing
  | otherwise = do
      -- resolve all args types
      targs <- fmap rights $ traverse compile args
      case sequence targs of
        Just t ->
          if length t /= length params
            then -- there was an error in one argument
              return $ Nothing
            else
              if all (\(a, b) -> a == b) $ zip t params -- compare types
                then return $ Nothing -- all good!
                else return $ Just ("type mismatch in function call\n  unexpected " ++ A.showList t ++ "\n   expecting " ++ A.showList params)
        Nothing ->
          -- there was an error in on argument
          return $ Nothing

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

-- | @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  unexpected ()\n   expecting " ++ showMaybet fret
typecheckReturn (Just value) fret = do
  r <- compile value
  case r of
    Right r ->
      if r == fret
        then return $ Nothing
        else return $ Just $ "invalid return value\n  unexpected " ++ showMaybet r ++ "\n   expecting " ++ showMaybet fret
    Left _ -> return $ Nothing -- error resolving return value

compile :: A.Expr -> State CompState CompResult
compile x = do
  case x of
    (A.Module name pos) -> return $ Right Nothing
    (A.Num _ _) -> return $ Right $ Just $ A.Type "u8" -- TODO: placeholder
    (A.BinOp _ a b) -> do
      l <- compile a
      r <- compile b
      return $ l -- TODO: placeholder
    (A.Func ident params ret body priv anon pos) -> do
      -- current env
      (ev, errs) <- get
      -- updated with the function
      (ev, errs) <-
        return $ case addSymUniq ev (ident, ftype, 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) params
      -- helper for return
      nev <- return $ addSym nev ("$fn$", ftype, pos)
      put (nev, errs)
      r <- compileAll body
      (_, errs) <- get
      -- store updated errors and the env with the function
      put (ev, errs)
      return $ Right $ Just ftype
      where
        ftype = A.toFuncType params ret
    (A.Call ident args pos) -> do
      r <- compile ident
      case r of
        Right (Just (A.FuncType params rtyp)) -> do
          r <- typecheckCall args params
          case r of
            Just err -> addError $ Error TypeError err pos
            Nothing -> return $ Right rtyp
        Right _ -> addError $ Error GenericError "non callable value in function call" pos
        _ -> return $ Right Nothing
    (A.Return value pos) -> do
      (ev, errs) <- get
      case getSyml ev "$fn$" of
        Nothing -> addError $ Error UnexpectedReturn "return without function call" pos
        Just (_, A.FuncType _ rtyp, _) -> do
          r <- typecheckReturn value rtyp
          case r of
            Just err -> addError $ Error TypeError err pos
            Nothing -> return $ Right rtyp
    (A.Var ident pos) -> do
      (ev, errs) <- get
      case getSym ev ident of
        Just (_, t, _) -> return $ Right $ Just t
        Nothing -> addError $ Error GenericError ("undefined variable \"" ++ ident ++ "\"") pos

compileAll :: [A.Expr] -> State CompState CompResult
compileAll (x : xs) = do
  compile x
  compileAll xs
compileAll [] = do
  (_, errs) <- get
  case errs of
    [] -> return $ Right Nothing
    _ -> return $ Left errs