module LambdaCube.SystemFw.TypeChecker where

import           Data.List               (uncons)
import           LambdaCube.SystemFw.Ast

substituteType :: Int -> LCType -> LCTerm -> LCTerm
substituteType n v = go n
  where
    go _ e@(LCVar _)  = e
    go m (LCLam t b)  = LCLam (substituteTypeInType m v t) $ go m b
    go m (LCApp f a)  = go m f `LCApp` go m a
    go m (LCTLam k b) = LCTLam k $ go (m + 1) b
    go m (LCTApp f t) = go m f `LCTApp` substituteTypeInType m v t

substituteTypeInType :: Int -> LCType -> LCType -> LCType
substituteTypeInType n v = go n
  where
    go _ LCBase        = LCBase
    go m e@(LCTVar l)  = if m == l then v else e
    go m (LCArr a b)   = go m a `LCArr` go m b
    go m (LCUniv k a)  = LCUniv k $ go (m + 1) a
    go m (LCTTLam k b) = LCTTLam k $ go (m + 1) b
    go m (LCTTApp f a) = go m f `LCTTApp` go m a

reduceType :: LCType -> LCType
reduceType = go
  where
    go LCBase = LCBase
    go e@(LCTVar _) = e
    go (LCArr a b) = go a `LCArr` go b
    go (LCUniv k a) = LCUniv k $ go a
    go (LCTTLam k b) = LCTTLam k $ go b
    go (LCTTApp f a)
      | LCTTLam _ b <- go f
      , v <- go a
      = go $ substituteTypeInType 0 v b
      | otherwise
      = error "Did you really kind check this?"

infer :: LCTerm -> Maybe LCType
infer = go []
  where
    go tl (LCVar n) = fmap fst . uncons $ drop n tl
    go tl (LCLam t b)
      | Just LCStar <- inferKind t
      = LCArr v <$> go (v : tl) b
      | otherwise
      = Nothing
      where
        v = reduceType t
    go tl (LCApp f a)
      | Just (LCArr at' rt) <- go tl f
      , Just at <- go tl a
      , at == at'
      = Just rt
      | otherwise
      = Nothing
    go tl (LCTLam k b) = LCUniv k <$> go tl b
    go tl (LCTApp f t)
      | Just (LCUniv tk' rt) <- go tl f
      , Just tk <- inferKind t
      , tk == tk'
      = Just $ substituteTypeInType 0 t rt
      | otherwise
      = Nothing

inferKind :: LCType -> Maybe LCKind
inferKind = go []
  where
    go _  LCBase = Just LCStar
    go kl (LCTVar n) = fmap fst . uncons $ drop n kl
    go kl (LCArr a b)
      | Just LCStar <- go kl a
      , Just LCStar <- go kl b
      = Just LCStar
      | otherwise
      = Nothing
    go kl (LCUniv k a) = go (k : kl) a
    go kl (LCTTLam k b) = LCKArr k <$> go (k : kl) b
    go kl (LCTTApp f a)
      | Just (LCKArr ak' rk) <- go kl f
      , Just ak <- go kl a
      , ak == ak'
      = Just rk
      | otherwise
      = Nothing