{-# LANGUAGE TupleSections, LambdaCase, OverloadedStrings #-}
module PostProcess where

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

import Data.Map (Map)
import Data.Set (Set)
import Data.Text (Text)

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

import qualified Hm.Abs as H
import Type
import Misc
import TC

import Prelude hiding (map)

-- Type env for parsing type signatures
type Process = StateT (Set Id) Check

insertType :: Id -> Process ()
insertType i = get >>= put . S.insert i

runProcess :: Process a -> Set Id -> Check a
runProcess = (fst <$>) .: runStateT

postprocess :: [H.Def] -> Process [TL]
postprocess defs = addDef <~> defs >> defToTL <~> defs

addDef :: H.Def -> Process ()
addDef = \case
  H.VarDef _ i t _ -> lift . addEnv i =<< typeSigToPolyT t -- the type checker will check this matches the exp
  H.TypeDef _ t _  -> insertType =<< fst <$> typeSigToIdParams t

-- add type before typesig to id params
defToTL :: H.Def -> Process TL
defToTL (H.VarDef p i t e) = VarDef p i <$> typeSigToPolyT t <*> expToExp e
defToTL (H.TypeDef p t ds) = do 
    (i,_) <- typeSigToIdParams t

    let (Id s) = i

    monoT <- declTo2pl <~> ds
    let env = M.fromList (map (\(i,m) -> (i, Mono m)) monoT)

    -- check that all constructors construct correct type
    mapM_ (guard (throwError (InvalidConstructor p)) . constructs i . snd) monoT

    -- check that there are no unbound variables
    guard (throwError (Unimplemented "Type parameters")) (free env == S.empty)

    -- add recursor to env
    tv <- lift fresh

    let replace = \case {
        TArr l r -> TArr (replace l) (replace r) ;
        TCon i'  -> if i' == i then tv else TCon i' ;
        TVar i'  -> TVar i' ;
      }

    recType <- lift . generalize . foldr TArr (tv `TArr` TCon i) . reverse . map (replace . snd) $ monoT

    let env' = M.insert (Id ("rec[" <> s <> "]")) recType env

    pure (TypeDef p i [] env')


typeSigToIdParams :: H.TypeSig -> Process (Id, [Id])
typeSigToIdParams = \case
  H.TypeFun p _ _ -> throwError (InvalidTypeDecl p)
  H.TypeApp {}    -> throwError (Unimplemented "Type parameters")
  H.TypeVar _ i   -> pure (i, [])

typeSigToPolyT :: H.TypeSig -> Process PolyT
typeSigToPolyT t = typeSigToMonoT t >>= lift . generalize

typeSigToMonoT :: H.TypeSig -> Process MonoT
typeSigToMonoT = \case
  H.TypeFun _ t1 t2 -> typeSigToMonoT t1 >>= (<$> typeSigToMonoT t2) . TArr
  H.TypeApp {}      -> throwError (Unimplemented "Type parameters")
  H.TypeVar _ i     -> get >>= \s -> pure $
                         if S.member i s
                            then TCon i
                            else TVar i

expToExp :: H.Exp -> Process Exp
expToExp = \case
  H.ExpLet p as e  -> Let p <$> assignTo2pl <~> as <*> expToExp e
  H.ExpTyped p e t -> Typed p <$> expToExp e <*> typeSigToMonoT t
  H.ExpAbs p is e  -> Abs p is <$> expToExp e
  H.ExpApp p e es  -> App p <$> expToExp e <*> expToExp <~> es
  H.ExpVar p i     -> pure (Var p i)

declTo2pl :: H.Decl -> Process (Id, MonoT)
declTo2pl (H.Decl _ i t) = (i,) <$> typeSigToMonoT t

assignTo2pl :: H.Assign -> Process (Id, Exp)
assignTo2pl (H.Assign _ i e) = (i,) <$> expToExp e