{-# LANGUAGE LambdaCase, TypeSynonymInstances #-}
{-# LANGUAGE TupleSections, FlexibleInstances #-}
module TC where

import Control.Monad.Reader hiding (guard)
import Control.Monad.State  hiding (guard)
import Control.Monad.Except hiding (guard)

import Data.Set (Set)
import qualified Data.Set as S
import qualified Data.Map as M

import qualified Data.Text as T

import Type
import Misc

import Prelude hiding (map)

map :: Functor f => (a -> b) -> f a -> f b
map = fmap

runCheck :: CheckState -> Check a -> (Either TypeError a, CheckState)
runCheck s = (flip runState) s . runExceptT . getCheck

-- I'm still not quite sure how replicateM works, but in this instance it is
-- used to generate a list of strings "a", "b" ... "z", "aa", "ab" ... so on
--
-- Does it make sense to start with an empty state?
initialState :: CheckState
initialState = CS ([1..] >>= map (Id . T.pack) . flip replicateM ['a'..'z']) Nothing M.empty

getVars :: Check [Id]
getVars = variables <$> get

setVars :: [Id] -> Check ()
setVars ids = get >>= \s -> put (CS ids (lastPos s) (typeEnv s))

getEnv :: Check TypeEnv
getEnv = typeEnv <$> get

setEnv :: TypeEnv -> Check ()
setEnv env = get >>= \s -> put (CS (variables s) (lastPos s) env)

addEnv :: Id -> PolyT -> Check ()
addEnv i p = getEnv >>= setEnv . M.insert i p

localEnv :: TypeEnv -> Check a -> Check a
localEnv e m = getEnv >>= \o -> setEnv e >> m >>= \r -> setEnv o >> pure r

localEnv' :: Check a -> Check a
localEnv' m = getEnv >>= \o -> m >>= \r -> setEnv o >> pure r

-- returns p again to allow chaining into lambdacase
setPos :: Positioned p => p -> Check p
setPos p = get >>= \s -> put (CS (variables s) (pos p) (typeEnv s)) >> pure p

guard :: Applicative f => f () -> Bool -> f ()
guard _ True = pure ()
guard f False = f

class Substitutable a where
  apply :: Subst -> a -> a -- ^ apply a substitution
  free   :: a -> Set Id    -- ^ free type variables

instance Substitutable MonoT where
  apply s = \case
    TCon i -> TCon i
    TVar i -> lookupDefault (TVar i) i s
    (t1 `TArr` t2) -> apply s t1 `TArr` apply s t2

  free = \case
    TCon{} -> S.empty
    TVar i -> S.singleton i
    (t1 `TArr` t2) -> free t1 <> free t2

instance Substitutable PolyT where
  apply s = \case
    Forall as t -> Forall as (apply (foldr M.delete s as) t)
    Mono t -> Mono (apply s t)

  free = \case
    Forall as t -> free t \\ as
    Mono t -> free t

instance Substitutable TypeEnv where
  apply s = map (apply s)
  free    = free . M.elems

instance Substitutable a => Substitutable [a] where
  apply = map . apply
  free  = foldMap free

applyEnv :: Subst -> Check ()
applyEnv s = getEnv >>=  setEnv . apply s

-- This substution, and that one
(<&>) :: Subst -> Subst -> Subst
(<&>) s1 s2 = map (apply s1) s2 <> s1

emptySubst :: Subst
emptySubst = M.empty

unify :: MonoT -> MonoT -> Check Subst
unify (l1 `TArr` r1) (l2 `TArr` r2) = do
  s1 <- unify l1 l2
  s2 <- unify (apply s1 r1) (apply s1 r2)
  pure (s1 <&> s2)

unify (TVar i) t = bind i t
unify t (TVar i) = bind i t

unify (TCon i1) (TCon i2) | i1 == i2 = pure emptySubst

unify t1 t2 = throwError (UnificationFailure t1 t2)

bind :: Id -> MonoT -> Check Subst
bind i1 (TVar i2) | i1 == i2            = pure emptySubst
bind i t          | S.member i (free t) = throwError (InfiniteType i t)
                  | otherwise           = pure (M.singleton i t)

fresh :: Check MonoT
fresh = do
  (var:vars) <- getVars
  setVars vars
  pure (TVar var)

-- replace polymorphic type variables with monomorphic ones
instantiate :: PolyT -> Check MonoT
instantiate (Mono t) = pure t
instantiate (Forall is t) = foldM freshInsert emptySubst is >>= pure . (flip apply) t
  where
    freshInsert :: Subst -> Id -> Check Subst
    freshInsert s k = (\a -> M.insert k a s) <$> fresh

generalize :: MonoT -> Check PolyT
generalize t = getEnv >>= \env -> pure (Forall (free t \\ free env) t)

lookupType :: Id -> Check MonoT
lookupType i = getEnv >>= \env ->
  case M.lookup i env of
    Nothing -> throwError (UnboundVariable i)
    Just t  -> instantiate t

constructs :: Id -> MonoT -> Bool
constructs i (TArr _ t) = constructs i t
constructs i1 (TCon i2) = i1 == i2
constructs _ _          = False

infer :: Exp -> Check (Subst, MonoT)
infer = setPos >=> \case
  Let   _ [] e            -> infer e
  Let   p ((i,e1):ies) e2 -> do
    (s1, t1) <- infer e1
    apply s1 <$> getEnv >>= \e -> localEnv e $ do
      t1g <- generalize t1
      addEnv i t1g
      (s2, t2) <- infer (Let p ies e2)
      pure (s2 <&> s1, t2)

  Abs   _ [] e     -> infer e
  Abs   p (i:is) e -> localEnv' $ do
    tv <- fresh
    addEnv i (Forall S.empty tv)
    (s, t) <- infer (Abs p is e)
    pure (s, apply s tv `TArr` t)

  App   p e1 e2 -> localEnv' $ do
    tv <- fresh
    (s1, t1) <- infer e1
    applyEnv s1
    (s2, t2) <- infer e2
    s3       <- unify (apply s2 t1) (TArr t2 tv)
    return (s3 <&> s2 <&> s1, apply s3 tv)
  App   _ _ _ -> throwError Oop

  Var   _ i   -> (emptySubst,) <$> lookupType i