{-# LANGUAGE TupleSections #-}

module Elaborator where

import Control.Monad.Except
import Data.List (elemIndex, lookup)
import Data.List.NonEmpty ((<|))
import Errors (Error (..), Result)
import Expr (Expr)
import qualified Expr as E
import IR (IRDef (..), IRExpr, IRProgram)
import qualified IR as I
import Relude.Extra.Lens

type Binders = [Text]

data SectionContext = SectionContext
    { sectionDefs :: [(Text, [Text])] -- name and list of variables it depends on
    , sectionVars :: [(Text, IRExpr)] -- variables and their types
    }

type ElabMonad = State SectionContext

lookupDefInCtxt :: Text -> SectionContext -> Maybe [(Text, IRExpr)]
lookupDefInCtxt def (SectionContext defs vars) = mapMaybe helper <$> lookup def defs
  where
    helper dep = (dep,) <$> dep `lookup` vars

-- looks up a definition in the context and gives a list of the variables and
-- their types that it depends on
lookupDef :: Text -> ElabMonad (Maybe [(Text, IRExpr)])
lookupDef def = lookupDefInCtxt def <$> get

lookupVarInCtxt :: Text -> SectionContext -> Maybe IRExpr
lookupVarInCtxt var = lookup var . sectionVars

-- looks up a variable in the context and returns its type
lookupVar :: Text -> ElabMonad (Maybe IRExpr)
lookupVar var = lookupVarInCtxt var <$> get

sectionDefsL :: Lens' SectionContext [(Text, [Text])]
sectionDefsL = lens sectionDefs setter
  where
    setter ctxt newDefs = ctxt{sectionDefs = newDefs}

sectionVarsL :: Lens' SectionContext [(Text, IRExpr)]
sectionVarsL = lens sectionVars setter
  where
    setter ctxt newVars = ctxt{sectionVars = newVars}

saveState :: ElabMonad a -> ElabMonad a
saveState action = get >>= (action <*) . put

elabSection :: Text -> [IRDef] -> ElabMonad [IRDef]
elabSection _name contents = saveState $ concat <$> forM contents elabDef

elabProgram :: IRProgram -> IRProgram
elabProgram prog = evalState (elabSection "" prog) (SectionContext [] [])

pushVariable :: Text -> IRExpr -> SectionContext -> SectionContext
pushVariable name ty (SectionContext defs vars) = SectionContext defs ((name, ty) : vars)

pushDefinition :: Text -> [Text] -> SectionContext -> SectionContext
pushDefinition name defVars (SectionContext defs vars) = SectionContext ((name, defVars) : defs) vars

removeName :: Text -> ElabMonad ()
removeName name = do
    modify $ over sectionDefsL (filter ((/= name) . fst))
    modify $ over sectionVarsL (filter ((/= name) . fst))

-- find all the section variables used in an expression
usedVars :: IRExpr -> ElabMonad [(Text, IRExpr)]
usedVars (I.Var name) = maybe [] (pure . (name,)) <$> lookupVar name
usedVars I.Star = pure []
usedVars (I.Level _) = pure []
usedVars (I.App m n) = (++) <$> usedVars m <*> usedVars n
usedVars (I.Abs name ty ascr body) = saveState $ do
    ty' <- usedVars ty
    ascr' <- traverse usedVars ascr
    removeName name
    ((ty' ++ (ascr' ?: [])) ++) <$> usedVars body
usedVars (I.Pi name ty ascr body) = saveState $ do
    ty' <- usedVars ty
    ascr' <- traverse usedVars ascr
    removeName name
    ((ty' ++ (ascr' ?: [])) ++) <$> usedVars body
usedVars (I.Let name ascr value body) = saveState $ do
    ty' <- usedVars value
    ascr' <- traverse usedVars ascr
    removeName name
    ((ty' ++ (ascr' ?: [])) ++) <$> usedVars body

-- traverse the body of a definition, adding the necessary section arguments to
-- any definitions made within the section
traverseBody :: IRExpr -> ElabMonad IRExpr
traverseBody (I.Var name) = do
    deps <- lookupDef name
    case deps of
        Nothing -> pure $ I.Var name
        Just [] -> pure $ I.Var name
        Just ((top, _) : rest) -> pure $ foldr (I.App . I.Var . fst) (I.Var top) rest
traverseBody I.Star = pure I.Star
traverseBody e@(I.Level _) = pure e
traverseBody (I.App m n) = I.App <$> traverseBody m <*> traverseBody n
traverseBody (I.Abs name ty ascr body) = saveState $ do
    ty' <- traverseBody ty
    ascr' <- traverse traverseBody ascr
    removeName name
    I.Abs name ty' ascr' <$> traverseBody body
traverseBody (I.Pi name ty ascr body) = saveState $ do
    ty' <- traverseBody ty
    ascr' <- traverse traverseBody ascr
    removeName name
    I.Pi name ty' ascr' <$> traverseBody body
traverseBody (I.Let name ascr value body) = saveState $ do
    value' <- traverseBody value
    ascr' <- traverse traverseBody ascr
    removeName name
    I.Let name ascr' value' <$> traverseBody body

mkPi :: (Text, IRExpr) -> IRExpr -> IRExpr
mkPi (param, typ) = I.Pi param typ Nothing

mkAbs :: (Text, IRExpr) -> IRExpr -> IRExpr
mkAbs (param, typ) = I.Abs param typ Nothing

generalizeType :: IRExpr -> [(Text, IRExpr)] -> IRExpr
generalizeType = foldr mkPi

generalizeVal :: IRExpr -> [(Text, IRExpr)] -> IRExpr
generalizeVal = foldr mkAbs

elabDef :: IRDef -> ElabMonad [IRDef]
elabDef (Def name ty body) = do
    tyVars <- fromMaybe [] <$> traverse usedVars ty
    bodyVars <- usedVars body
    let totalVars = tyVars ++ bodyVars
    modify $ pushDefinition name (map fst totalVars)
    pure [Def name (flip generalizeType totalVars <$> ty) (generalizeVal body totalVars)]
elabDef (Axiom name ty) = do
    vars <- usedVars ty
    modify $ pushDefinition name (map fst vars)
    pure [Axiom name (generalizeType ty vars)]
elabDef (Section name contents) = saveState $ elabSection name contents
elabDef (Variable name ty) = [] <$ modify' (pushVariable name ty)

-- saveBinders :: State Binders a -> State Binders a
-- saveBinders action = do
--     binders <- get
--     res <- action
--     put binders
--     pure res
--
-- elaborate :: IRExpr -> Expr
-- elaborate ir = evalState (elaborate' ir) []
--   where
--     elaborate' :: IRExpr -> State Binders Expr
--     elaborate' (I.Var n) = do
--         binders <- get
--         pure $ E.Var n . fromIntegral <$> elemIndex n binders ?: E.Free n
--     elaborate' I.Star = pure E.Star
--     elaborate' (I.Level level) = pure $ E.Level level
--     elaborate' (I.App m n) = E.App <$> elaborate' m <*> elaborate' n
--     elaborate' (I.Abs x t a b) = saveBinders $ do
--         t' <- elaborate' t
--         a' <- traverse elaborate' a
--         modify (x :)
--         E.Abs x t' a' <$> elaborate' b
--     elaborate' (I.Pi x t a b) = saveBinders $ do
--         t' <- elaborate' t
--         a' <- traverse elaborate' a
--         modify (x :)
--         E.Pi x t' a' <$> elaborate' b
--     elaborate' (I.Let name Nothing val body) = saveBinders $ do
--         val' <- elaborate' val
--         modify (name :)
--         E.Let name Nothing val' <$> elaborate' body
--     elaborate' (I.Let name (Just ty) val body) = saveBinders $ do
--         val' <- elaborate' val
--         ty' <- elaborate' ty
--         modify (name :)
--         E.Let name (Just ty') val' <$> elaborate' body