{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE DeriveAnyClass #-}

module Minilang.Lambda.Infer where

import Control.Monad.State (MonadState (..), StateT, evalStateT, modify)
import Data.Aeson (FromJSON, ToJSON)
import Data.Bifunctor (first)
import Data.Char (chr, ord)
import qualified Data.List as List
import qualified Data.Map as Map
import Data.Text (Text, pack)
import GHC.Generics (Generic)
import qualified Minilang.Lambda.Eval as Eval
import Minilang.Lambda.Unify (Type (..), UnifyError, apply, unify)
import Prelude hiding (lookup)
import qualified Minilang.Lambda.Unify as Unify

data ATerm ann
  = Var Text ann
  | Lam Text (ATerm ann) ann
  | App (ATerm ann) (ATerm ann) ann
  deriving (Show, Eq)

data TypeError
  = UnifyError UnifyError
  deriving (Show, Eq)

render :: TypeError -> Text
render (UnifyError ue) = Unify.render ue

fromTerm :: Eval.Term -> ATerm ()
fromTerm (Eval.Var x)     = Var x ()
fromTerm (Eval.Lam x body) = Lam x (fromTerm body) ()
fromTerm (Eval.App a b)   = App (fromTerm a) (fromTerm b) ()
-- Literal values (Num, Str) are treated as typed by a fresh variable;
-- the inferencer only needs to traverse the term structure, not inspect
-- literal payloads.
fromTerm (Eval.Num _)     = Var "<num>" ()
fromTerm (Eval.Str _)     = Var "<str>" ()

infer :: Eval.Term -> Either TypeError Type
infer t = do
  annotated <- annotate t
  constraints <- collect [annotated] []
  subs <- first UnifyError $ unify constraints
  pure $ apply subs (typeOf annotated)

collect :: [ATerm Type] -> [(Type, Type)] -> Either TypeError [(Type, Type)]
collect terms constraints = case terms of
  [] -> Right constraints
  (Var _ _) : rest -> collect rest constraints
  (Lam _ body _) : rest -> collect (body : rest) constraints
  (App x y t) : rest ->
    let (f, b) = (typeOf x, typeOf y)
     in collect (x : y : rest) ((f, b :-> t) : constraints)

typeOf :: ATerm Type -> Type
typeOf (Var _ t) = t
typeOf (Lam _ _ t) = t
typeOf (App _ _ t) = t

data Context = Context
  { nextVar :: Int,
    typeContext :: Map.Map Text Type
  }
  deriving (Eq, Show, Generic, ToJSON, FromJSON)

extendContext :: Text -> Type -> Context -> Context
extendContext x ty (Context n env) =
  Context n (Map.insert x ty env)

annotate :: Eval.Term -> Either TypeError (ATerm Type)
annotate t =
  evalStateT (go (fromTerm t) []) newContext
  where
    go :: ATerm a -> [(Text, Type)] -> StateT Context (Either TypeError) (ATerm Type)
    go t bounds = case t of
      Var x _ -> case List.lookup x bounds of
        Just ty -> pure $ Var x ty
        Nothing ->
          lookup x >>= \case
            Just ty -> pure $ Var x ty
            Nothing -> do
              v <- freshVar
              bind x v
              pure $ Var x v
      Lam x body _ -> do
        ty <- freshVar
        body' <- go body ((x, ty) : bounds)
        pure $ Lam x body' (ty :-> typeOf body')
      App a b _ ->
        App <$> go a bounds <*> go b bounds <*> freshVar

newContext :: Context
newContext = Context 0 Map.empty

bind :: Text -> Type -> StateT Context (Either TypeError) ()
bind x ty = do
  modify (\(Context n env) -> Context n (Map.insert x ty env))

lookup :: Text -> StateT Context (Either TypeError) (Maybe Type)
lookup x = do
  Context _ env <- get
  pure $ Map.lookup x env

freshVar :: (Monad m) => StateT Context m Type
freshVar = do
  get >>= \(Context n _) ->
    modify (\(Context n' env') -> Context (n' + 1) env') >> pure (TyVar $ nameOf n)
  where
    nameOf n
      | n < 26 = pack [chr (n + ord 'a')]
      | otherwise =
          let (d, r) = n `divMod` 26
           in pack $ chr (r + ord 'a') : show d