{-# LANGUAGE DuplicateRecordFields, OverloadedLists, InstanceSigs #-}
{-# OPTIONS_GHC -Wall #-}
module Rubyhs.References
  ( References(entries)
  , Entry(..)
  , FQN(..)
  , references
  , Env(..)
  , Result(..)
  , Namespace(..)
  , Context(..)
  ) 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.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

data Context = ContextFun FQN | ContextMod Namespace

deriving stock instance Eq Context
deriving stock instance Ord Context

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

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

data Result = Result
  { declarations :: Set FQN
  , applications :: Map Context (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) = (prettyContext x, y)

prettyContext :: Context -> Text
prettyContext = \case
  ContextFun q -> prettyFQN q
  ContextMod ns -> prettyNamespace ns
  
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 => (Context -> Context) -> 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 (ContextMod q) = ContextMod $ name2ns n <> q
  go _ = error "Cannot append module to context in function context"

name2ns :: Name -> Namespace
name2ns (Name o) = go o
  where
  go :: Aeson.Value -> Namespace
  go = \case
    Aeson.Array [Aeson.String "const", x, Aeson.String s] -> case x of
      Aeson.Null -> [s]
      _ -> go x <> [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) emptyEnv
  emptyEnv = Env mempty mempty (ContextMod mempty)

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

instance References Statement where
  entries = \case
    StmtModule m -> entries m
    StmtFunction f -> entries f
    StmtSend s -> entries s
    -- TODO:
    StmtRBlock{} -> pure ()
    StmtAnything{} -> pure ()

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

instance References Function where
  entries Function{name, block} = do
    namespace <- getContext >>= \case
      ContextMod c -> pure c
      _ -> error "Cannot have a function declaration in a function context"
    declaration $ qual namespace name
    locally $ do
      writeContext (ContextFun $ 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