{-# LANGUAGE DuplicateRecordFields, OverloadedLists, InstanceSigs #-}
{-# OPTIONS_GHC -Wall #-}
module Rubyhs.References
  ( References(entries)
  , FQN(..)
  , references
  , Env(..)
  , Result(..)
  , Context(..)
  , graph
  , prettyContext
  , Node(..)
  ) where

import Frelude
-- import Data.Language.Ruby hiding (context)
import qualified Data.Language.Ruby as Ruby
import Data.HashMap.Strict (HashMap)
import qualified Data.HashMap.Strict as HashMap
import qualified Data.Aeson.Encoding.Internal as Aeson
import Data.Kind
import Control.Monad.State
import qualified Data.Aeson as Aeson
import qualified Data.Aeson.Types as Aeson
import Data.Set (Set)
import qualified Data.Set as Set
import qualified Data.Text as Text
import Data.Coerce
import Data.Graph (Graph, Vertex, Forest)
import qualified Data.Graph as Graph
import Data.Hashable
import qualified Data.Char as Char

span :: G node key -> Forest node
span (g, f, _) = fmap ((\(x, _, _) -> x) . f) <$> Graph.dff g

type G node key = (Graph, Vertex -> (node, key, [key]), key -> Maybe Vertex)

graph :: Ruby.Begin -> Forest Node
graph b = span $ Graph.graphFromEdges $ go <$> toList ys
  where
  go :: (Node, Set Node) -> (Node, Text, [Text])
  go (x, xs) = (x, prettyContext x, prettyContext <$> toList xs)
  Result ys = references b

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 = coerce
  toList = coerce
deriving newtype instance Semigroup Namespace
deriving newtype instance Monoid Namespace
instance FromJSON Namespace where
  parseJSON = Aeson.withText "Namespace" $ \t -> Namespace <$> (traverse pure $ Text.splitOn "::" t)
instance ToJSON Namespace where
  toJSON = Aeson.String . showNamespace

showNamespace :: Namespace -> Text
showNamespace = Text.intercalate "::" . coerce

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

data Node = NodeDef FQN | NodeModule Namespace

beginsWithCapital :: Text -> Bool
beginsWithCapital = Char.isUpper . Text.head

deriving stock instance Show Node
deriving stock instance Eq Node
deriving stock instance Ord Node
instance FromJSON Node where
  parseJSON = Aeson.withText "Node" $ nodeParser
nodeParser :: Text -> Aeson.Parser Node
nodeParser = \t -> case Text.splitOn "." t of
    [mods,fun] -> fmap NodeDef (FQN <$> Aeson.parseJSON (Aeson.String mods) <*> pure fun)
    [mods]     ->
      if beginsWithCapital mods
      then NodeModule <$> Aeson.parseJSON (Aeson.String mods)
      else pure $ NodeDef $ FQN mempty mods
    _          -> empty
instance Aeson.FromJSONKey Node where
  fromJSONKey = Aeson.FromJSONKeyTextParser nodeParser
instance Aeson.ToJSONKey Node where
  toJSONKey = Aeson.ToJSONKeyText prettyContext (\n -> Aeson.text $ prettyContext n)
instance ToJSON Node where
  toJSON = \case
    NodeDef q -> Aeson.toJSON q
    NodeModule m -> Aeson.toJSON m
instance Hashable Node where
  hashWithSalt n a = hashWithSalt @Text n $ prettyContext a

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

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{name, namespace} = case namespace of
  [] -> name
  _ -> convertString (showNamespace namespace) <> "." <> name

newtype Context = Context Node

deriving stock instance Eq Context
deriving stock instance Ord Context
deriving newtype instance ToJSON Context
-- HACK: Not really a semigroup as is evident from the implementation.
instance Semigroup Context where
  Context c0 <> Context c1 = case (c0, c1) of
    (NodeModule n0, NodeModule n1)      -> Context $ NodeModule $ n0 <> n1
    (NodeModule n0, NodeDef (FQN n1 f)) -> Context $ NodeDef $ FQN (n0 <> n1) f
    (NodeDef{}, NodeModule{})           -> error "Cannot append module to function context."
    (NodeDef{}, NodeDef{})              -> error "Cannot append function to function context."

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

data Env = Env
  { declarations :: Set Node
  , applications :: HashMap Node (Set Node)
  , context      :: Context
  }

newtype Result = Result (HashMap Node (Set Node))
-- HACK: FromJSON and ToJSON don't agree.
deriving newtype instance FromJSON Result
instance ToJSON Result where
  toJSON (Result a) = Aeson.Object $ fromList $ go <$> toList a
    where
    go :: (Node, Set Node) -> (Text, Aeson.Value)
    go (x, y) = (prettyContext x, Aeson.toJSON y)

prettyContext :: Node -> Text
prettyContext = \case
  NodeDef fun -> prettyFQN fun
  NodeModule ns -> showNamespace 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 n = do
    Context c <- getContext
    let
      go env@Env{applications}
        = env { applications = HashMap.insertWith mappend c (Set.singleton n) applications }
    modify go
  getContext = gets Rubyhs.References.context
  writeContext q = modify go
    where
    go env = env { Rubyhs.References.context = q }

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

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

references :: Ruby.Begin -> Result
references q = Result $ HashMap.unionWith mappend applications declarations'
  where
  Env{declarations,applications} = execState (entries @_ @(State Env) q) emptyEnv
  emptyEnv = Env mempty mempty (Context $ NodeModule mempty)
  declarations' = fromList $ (\x -> (x, mempty)) <$> toList declarations

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

instance References Ruby.Statement where
  entries = \case
    Ruby.StmtModule m   -> entries m
    Ruby.StmtDef f      -> entries f
    Ruby.StmtDefs f     -> entries f
    Ruby.StmtSend s     -> entries s
    Ruby.StmtConst c    -> entries c
    Ruby.StmtBlock b    -> entries b
    Ruby.StmtCasgn c    -> entries c
    Ruby.StmtArray a    -> entries a
    Ruby.StmtSym s      -> entries s
    Ruby.StmtStr s      -> entries s
    Ruby.StmtLvasgn a   -> entries a
    Ruby.StmtLvar a     -> entries a
    Ruby.StmtIvar a     -> entries a
    Ruby.StmtSelf s     -> entries s
    Ruby.StmtCbase s    -> entries s
    Ruby.StmtNil n      -> entries n
    Ruby.StmtAnything a -> entries a

instance References Ruby.Block where
  entries Ruby.Block{send,args,begin} = do
    entries send
    entries args
    entries begin
instance References Ruby.Args where
  entries = const $ pure ()
-- TODO: We have to make a "declaration" for the constant here as
-- well!
instance References Ruby.Casgn where
  entries Ruby.Casgn{rhs} = entries rhs
instance References Ruby.Sym where
  entries _ = pure ()
instance References Ruby.Str where
  entries _ = pure ()
instance References Ruby.Lvasgn where
  entries _ = pure ()
instance References Ruby.Lvar where
  entries _ = pure ()
instance References Ruby.Ivar where
  entries _ = pure ()
instance References Ruby.Self where
  entries _ = pure ()
instance References Ruby.Cbase where
  entries _ = pure ()
instance References Ruby.Nil where
  entries _ = pure ()
instance References Ruby.Array where
  entries Ruby.Array{statements} = traverse_ entries statements
instance References Ruby.Anything where
  entries = const $ pure ()

instance References Ruby.Const where
  entries con = application $ NodeModule $ constToNamespace con

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

instance References Ruby.Module where
  entries Ruby.Module{name, begin} = do
    updateContext $ (<>) $ Context $ NodeModule $ constToNamespace name
    c <- getContext >>= \case
      Context (NodeModule c) -> pure c
      _ -> error "..."
    declaration $ NodeModule c
    entries begin

atomToNode :: MyMonad m => Ruby.Atom -> m Node
atomToNode (Ruby.Atom name) = do
  namespace <- getContext >>= \case
    Context (NodeModule c) -> pure c
    -- We could allow this and just say that the function defined in
    -- another function sits in the same context as the surrounding
    -- function.
    Context NodeDef{} -> error "Cannot have a function declaration in a function context"
  pure $ NodeDef $ FQN namespace name

instance References Ruby.Def where
  entries Ruby.Def{atom, begin} = do
    node <- atomToNode atom
    declaration node
    locally $ do
      writeContext (Context node)
      entries begin

instance References Ruby.Defs where
  -- TODO: The field `context` is also relevant here!
  entries Ruby.Defs{atom, begin}
    = entries $ Ruby.Def{atom,begin,args=error "hack"}

qual :: Namespace -> Ruby.Atom -> FQN
qual namespace (Ruby.Atom name) = FQN { namespace , name }

instance References Ruby.Send where
  entries Ruby.Send{context, atom, args} = do
    application $ NodeDef $ qual (statementToNamespace context) atom
    traverse_ entries args

statementToNamespace :: Ruby.Statement -> Namespace
statementToNamespace = go mempty
  where
  go acc = \case
    Ruby.StmtConst c    -> acc <> constToNamespace c
    -- The nil-case and cbase-case should produce different results,
    -- surely.  `Namespace` may not be a good representation.
    Ruby.StmtNil{}      -> acc
    Ruby.StmtCbase{}    -> acc
    -- The send-, ivar- and lvar- case cannot be handled because of the way
    -- we've defined `Namespace`.
    Ruby.StmtSend{}     -> acc
    Ruby.StmtIvar{}     -> acc
    Ruby.StmtLvar{}     -> acc
    _              -> error "Can only build namespaces from sequences of `const` statements"

constToNamespace :: Ruby.Const -> Namespace
constToNamespace Ruby.Const{context, atom} = statementToNamespace context <> [k]
  where
  Ruby.Atom k = atom