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

import Frelude
import Data.Language.Ruby hiding (Namespace)
import qualified Data.Language.Ruby as Ruby
import Data.Map (Map)
import qualified Data.Map as Map
import Data.Kind
import Control.Monad.State
import Data.Aeson ((.=))
import qualified Data.Aeson as Aeson
import Data.List
import Data.HashMap.Strict (HashMap)
import Data.Text.Prettyprint.Doc.Render.Text (renderStrict)
import Data.Text.Prettyprint.Doc (layoutCompact)
import Data.Set (Set)
import qualified Data.Set as Set

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

newtype Namespace = Namespace [Text]
deriving stock instance Show Namespace
deriving stock instance Eq Namespace
deriving stock instance Ord Namespace
instance IsList Namespace where
  type Item Namespace = Text
  fromList = Namespace . fromList
  toList (Namespace l) = toList l
deriving newtype instance Semigroup Namespace
deriving newtype instance Monoid Namespace

-- Names are in reverse order.
data FQN = FQN
  { namespace :: Namespace
  , name :: Text
  }

deriving stock instance Show FQN
deriving stock instance Eq FQN
deriving stock instance Ord FQN
instance Pretty Namespace where
  pretty = pretty . prettyNamespace

instance Aeson.ToJSON FQN where
  toJSON = Aeson.String . prettyFQN
instance Aeson.ToJSONKey FQN where
instance Pretty FQN where
  pretty = pretty . prettyFQN

prettyFQN :: FQN -> Text
prettyFQN FQN{namespace, name} = case namespace of
  [] -> name
  _ -> prettyNamespace namespace <> "." <> name

prettyNamespace :: Namespace -> Text
prettyNamespace = convertString . intercalate "::" . fmap convertString . reverse . toList

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

data Env = Env
  { declarations :: Set FQN
  , applications :: Map Namespace (Set FQN)
  , context      :: Namespace
  }

data Result = Result
  { declarations :: Set FQN
  , applications :: Map Namespace (Set FQN)
  }
instance ToJSON Result where
  toJSON Result{declarations,applications} = Aeson.object
    [ "declarations" .= declarations
    , "applications" .= f applications
    ]
    where
      f x = fromList @(HashMap _ _) $ go <$> toList x
      go (x, y) = (renderStrict $ layoutCompact $ pretty x, y)

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 = modify go
    where
    go :: Env -> Env
    go env@Env{declarations} = env { declarations = Set.insert q declarations }
  application q n = modify go
    where
    go :: Env -> Env
    go env@Env{applications} = env { applications = Map.insertWith mappend q (Set.singleton n) applications }
  getContext = gets context
  writeContext q = modify go
    where
    go env = env { context = q }

updateContext :: MyMonad m => (Namespace -> Namespace) -> 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 q = name2ns n <> q

name2ns :: Name -> Namespace
name2ns (Name o) = case o of
  Aeson.Array [_, _, Aeson.String s] -> [s]
  _ -> error $ show o

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

references :: Block -> Result
references q = Result{declarations, applications}
  where
  Env{declarations,applications} = 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
    namespace <- getContext
    declaration $ qual namespace name
    entries block

qual :: Namespace -> Name -> FQN
qual namespace (Name o) = case o of
  Aeson.String name -> FQN { namespace , name }
  _ -> error $ show o

instance References Send where
  entries Send{namespace, name} = do
    c <- getContext
    application c $ qual (fromNS namespace) name
    where
    fromNS :: Ruby.Namespace -> Namespace
    fromNS (Ruby.Namespace l) = Namespace $ go <$> l
    go :: Name -> Text
    go (Name o) = case o of
      Aeson.String s -> s
      _ -> error $ show o