{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE MultiParamTypeClasses #-}

module Check where

import           AST
import           Printer              ()
import           Reduction

import           Control.Monad.Except
import           Control.Monad.Reader

import           Text.Megaparsec      (SourcePos)

type Context m v i = [Term m v i]

type Judgement m v i a = ReaderT (Context m v i) (Either String) a

typeSynth :: (Show v, Enum i) => Term SourcePos v i -> Judgement SourcePos v i (PseudoTerm SourcePos v i)
typeSynth term' =
  case term' of
    Type metadata -> return $ Kind' (Kind metadata)
    Variable variable index position -> do
      context <- ask
      if Prelude.length context <= fromEnum index - 1
        then throwError $ "invalid variable encountered during type synthesis: "
                          ++ show variable ++ "\nat " ++ show position
        else return $ Term $ context !! (fromEnum index - 1)
    Application position argument function -> do
      function' <- typeSynth function
      case function' of
        Term f -> return $ Term $ Application position argument f
        Kind' _ -> throwError $ "invalid kind constructed during type synthesis"
                                ++ "\nat " ++ show position
    Abstraction variable position parameter body -> do
      body' <- local (parameter :) (typeSynth body)
      case body' of
        Term b  -> return $ Term $ Abstraction variable position parameter b
        Kind' b -> return $ Kind' $ Function variable position parameter b

fTypeSynth :: (Show v, Enum i) => Term SourcePos v i -> Judgement SourcePos v i (PseudoTerm SourcePos v i)
fTypeSynth t = do
  typ <- typeSynth t
  case typ of
    Term t' -> fTypeSynth t'
    Kind' _ -> return typ

termCorrect :: (Enum i, Show v, BetaEq (PseudoTerm SourcePos) v i) => Term SourcePos v i -> Judgement SourcePos v i ()
termCorrect t =
  case t of
    Type _ -> return ()
    Variable{} ->
      void (typeSynth t)
    Abstraction _ _ parameter body -> do
      termCorrect parameter
      local (parameter :) (termCorrect body)
    Application pos argument function -> do
      argType <- typeSynth argument
      termCorrect argument
      Term funType <- fTypeSynth function
      case nfTerm funType of
        Abstraction _ _ parameter body ->
          if Term parameter === argType
            then termCorrect (nfTerm $ substituteTerm (toEnum 1) argument body)
            else let argPosition = show $ getNodeMetadata (Term argument)
                     parPosition = show $ getNodeMetadata (Term parameter)
                 in throwError $ "argument type and function parameter do not match: \n"
                                 ++ "at " ++
                                    argPosition ++
                                    " with argument type of:   " ++
                                    show argType ++
                                    "\n" ++
                                    "at " ++
                                    parPosition ++
                                    " with parameter value of: " ++
                                    show parameter ++
                                    "\n"
        e@_ -> throwError $ "attempted application onto non function type at "
                            ++ show pos ++ "\nwith type: " ++ show e