{-# LANGUAGE UndecidableInstances #-} -- Eq/Ord/Show deriving

module Momo.Env
  ( Env
  , fromSignature
  , addSignature

  , addValue
  , findValue
  , addType
  , findType
  , addModule
  , findModule

  , Binding(..)
  , addSpec

  , find
  , FindError(..)
  ) where

import Control.Monad.Catch (Exception, MonadThrow(throwM))
import Data.Map.Strict qualified as Map
import Data.Text (Text)

import Momo.CoreSyntax qualified as Core
import Momo.Ident (Ident)
import Momo.Ident qualified as Ident
import Momo.ModSyntax qualified as Mod
import Momo.Path (Path)
import Momo.Path qualified as Path
import Momo.Subst (Subst)

data Binding term
  = Value (Core.Val term)
  | Type (Mod.TypeDecl term)
  | Module (Mod.ModType term)

deriving instance Core.EqTerm term => Eq (Binding term)
deriving instance Core.OrdTerm term => Ord (Binding term)
deriving instance Core.ShowTerm term => Show (Binding term)

type Env term = Ident.Table (Binding term)

addSignature :: [Mod.Specification term] -> Env term -> Env term
addSignature specs env = foldr addSpec env specs

fromSignature :: [Mod.Specification term] -> Env term
fromSignature specs = addSignature specs mempty

addSpec
  :: Mod.Specification term
  -> Env term
  -> Env term
addSpec = \case
  Mod.ValueSig  ident vty  -> addValue ident vty
  Mod.TypeSig   ident decl -> addType ident decl
  Mod.ModuleSig ident mty  -> addModule ident mty

addValue :: Ident -> Core.Val term -> Env term -> Env term
addValue ident = Map.insert ident . Value

addType :: Ident -> Mod.TypeDecl term -> Env term -> Env term
addType ident = Map.insert ident . Type

addModule :: Ident -> Mod.ModType term -> Env term -> Env term
addModule ident = Map.insert ident . Module

data FindError
  = ValueFieldExpected Path
  | TypeFieldExpected Path
  | ModuleFieldExpected Path
  | BindingNotFound Ident
  | StructureExpected Path
  | StructureFieldNotFound Path Text
  deriving (Eq, Ord, Show)

instance Exception FindError

findValue
  :: ( Core.CoreSyntax term
     , MonadThrow m
     )
  => Path
  -> Env term
  -> m (Core.Val term)
findValue path env =
  find path env >>= \case
    Value val ->
      pure val
    _ ->
      throwM $ ValueFieldExpected path

findType
  :: ( Core.CoreSyntax term
     , MonadThrow m
     )
  => Path
  -> Env term
  -> m (Mod.TypeDecl term)
findType path env =
  find path env >>= \case
    Type decl ->
      pure decl
    _ ->
      throwM $ TypeFieldExpected path

findModule
  :: ( Core.CoreSyntax term
     , MonadThrow m
     )
  => Path
  -> Env term
  -> m (Mod.ModType term)
findModule path env =
  find path env >>= \case
    Module mty ->
      pure mty
    _ ->
      throwM $ ModuleFieldExpected path

find
  :: forall term m
  . ( Core.CoreSyntax term
    , MonadThrow m
    )
  => Path
  -> Env term
  -> m (Binding term)
find queryPath env =
  case queryPath of
    Path.Ident{ident} ->
      case Map.lookup ident env of
        Nothing ->
          throwM $ BindingNotFound ident
        Just binding ->
          pure binding

    Path.Dot{path, field} ->
      findModule @term path env >>= \case
        Mod.Signature sg ->
          findField path field mempty sg
        Mod.FunctorType{} ->
          -- XXX: non-instantiated functor
          throwM $ StructureExpected path

findField
  :: forall term m
  .  ( Core.CoreSyntax term
     , MonadThrow m
     )
  => Path
  -> Text
  -> Subst
  -> [Mod.Specification term]
  -> m (Binding term)
findField path field subst = \case
  [] ->
    throwM $ StructureFieldNotFound path field

  Mod.ValueSig ident vty : remaining ->
    if ident.name == field then
      pure $ Value (Core.substVal @term vty subst)
    else
      findField path field subst remaining

  Mod.TypeSig ident decl : remaining ->
    if ident.name == field then
      pure $ Type (Mod.substTypeDecl decl subst)
    else
      let
        subst' =
          Map.insert
            ident
            (Path.Dot path ident.name)
            subst
      in
        findField path field subst' remaining

  Mod.ModuleSig ident mty : remaining ->
    if ident.name == field then
      pure $ Module (Mod.substModType mty subst)
    else
      let
        subst' =
          Map.insert
            ident
            (Path.Dot path ident.name)
            subst
      in
        findField path field subst' remaining