module Elaborator where

import Control.Monad.Except
import Data.List (elemIndex)
import Data.List.NonEmpty ((<|))
import Errors (Error (..), Result)
import Expr (Expr)
import qualified Expr as E
import IR (IRDef (..), IRExpr)
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 = StateT (NonEmpty SectionContext) Result

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 $ do
    modify (SectionContext [] [] <|)
    concat <$> forM contents elabDef

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

elabDef :: IRDef -> ElabMonad [IRDef]
elabDef (Def name ty body) = undefined
elabDef (Axiom name ty) = undefined
elabDef (Section name contents) = elabSection name contents
elabDef (Variable name ty) = [] <$ modify' (pushVariable name ty)

traverseDef :: [Text] -> IRExpr -> ElabMonad [Text]
traverseDef seen (I.Var _) = pure seen
traverseDef seen I.Star = pure seen
traverseDef seen (I.Level _) = pure seen
traverseDef seen (I.App m n) = (++) <$> traverseDef seen m <*> traverseDef seen n

-- game plan:
-- 1. add arguments to section local definitions

-- 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
--     helper :: (Monad m) => Maybe a -> (a -> m b) -> m (Maybe b)
--     helper Nothing _ = pure Nothing
--     helper (Just x) f = Just <$> f x
--
--     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' <- helper a elaborate'
--         modify (x :)
--         E.Abs x t' a' <$> elaborate' b
--     elaborate' (I.Pi x t a b) = saveBinders $ do
--         t' <- elaborate' t
--         a' <- helper a elaborate'
--         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