{-# LANGUAGE DuplicateRecordFields, OverloadedLists, InstanceSigs, RecordWildCards #-}
{-# OPTIONS_GHC -Wall #-}
module Rubyhs.References
  ( References(entries)
  , Entry(..)
  , FQN(..)
  , references
  ) where

import Frelude
import Data.List.NonEmpty (NonEmpty)
import Data.Language.Ruby
import Data.Map
import qualified Data.Map as Map
import Data.Kind
import Control.Monad.State
import qualified Data.Aeson as Aeson

data Entry a = Entry
  { node :: a
  , fqn :: FQN
  }

newtype FQN = FQN (NonEmpty Name)

deriving newtype instance Semigroup FQN
instance Monoid FQN where
  mempty = FQN $ pure $ Name Aeson.Null

deriving newtype instance Eq FQN
deriving newtype instance Ord FQN

class Monad m => MyMonad (m :: Type -> Type) where
  declaration :: FQN -> Name -> m ()
  application :: FQN -> Name -> m ()
  getContext :: m FQN
  writeContext :: FQN -> m ()

data Env = Env
  { declarations :: Map FQN [Name]
  , applications :: Map FQN [Name]
  , context      :: FQN
  }

instance Semigroup Env where
  Env a0 a1 a2 <> Env b0 b1 b2 = Env (a0 <> b0) (a1 <> b1) (a2 <> b2)

instance Monoid Env where
  mempty = Env mempty mempty mempty
  
instance MyMonad (State Env) where
  declaration q n = modify go
    where
    go env@Env{declarations} = env { declarations = Map.insertWith mappend q (pure n) declarations }
  application q n = modify go
    where
    go env@Env{applications} = env { applications = Map.insertWith mappend q (pure n) applications }
  getContext = gets context
  writeContext q = modify go
    where
    go env = env { context = q }

updateContext :: MyMonad m => (FQN -> FQN) -> m ()
updateContext f = getContext >>= \c -> writeContext (f c)

locally :: MyMonad m => m a -> m a
locally act = do
  old <- getContext
  res <- act
  writeContext old
  pure res

appendToContext :: MyMonad m => Name -> m ()
appendToContext n = updateContext go
  where
  go (FQN ns) = FQN $ pure n <> ns

class References a where
  entries :: MyMonad m => a -> m ()

references :: Block -> Env
references q = execState (entries @_ @(State Env) q) mempty

instance References Block where
  entries :: forall m . MyMonad m => Block -> m ()
  entries (Block defs) = traverse_ (locally . entries) defs

instance References Definition where
  entries = \case
    DefModule m -> entries m
    DefFunction f -> entries f
    DefSend s -> entries s

instance References Module where
  entries Module{name, block} = do
    appendToContext name
    entries block

instance References Function where
  entries Function{name, block} = do
    c <- getContext
    declaration c name
    entries block

instance References Send where
  entries Send{name} = do
    c <- getContext
    application c name