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

import qualified Prelude
import Frelude
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

-- The elements appear in reverse order.  I expect this to cause all
-- sorts of issues.
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

-- 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)
    -- The top-level
    [[]]       -> pure $ NodeModule mempty
    [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

prettyContext :: Node -> Text
prettyContext = \case
  NodeDef fun -> prettyFQN fun
  NodeModule ns -> showNamespace ns

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

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
  -- I think that the order must be swapped because they modules are in reverse order.
  Context c0 <> Context c1 = Context $ c0 <> c1
instance Semigroup Node where
  c0 <> c1 = case (c0, c1) of
    -- (NodeModule n0, NodeModule n1)      -> Context $ NodeModule $ n0 <> n1
    (NodeModule n0, NodeModule n1)      -> NodeModule $ n0 <> n1
    -- (NodeModule n0, NodeDef (FQN n1 f)) -> Context $ NodeDef $ FQN (n0 <> n1) f
    (NodeModule n0, NodeDef (FQN n1 f)) -> NodeDef $ FQN (n0 <> n1) f
    -- TODO: Problematic:  Is a Node really a semigroup with this representation?
    (NodeDef{}, NodeModule{})           -> c1
    (NodeDef (FQN n0 _), NodeDef (FQN n1 a1)) -> NodeDef $ FQN (n0 <> n1) a1

class Monad m => ReferenceM (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)
  
instance ReferenceM (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
    decls <- gets declarations
    n' <- findClosest decls n
    let
      -- n' = findCloses n
      -- case n of
      --   -- TODO: Maybe check if there is `a` is a function in the
      --   -- current closure.
      --   (NodeDef (FQN [] a)) -> NodeDef $ FQN (nodeNs c) a
      --   -- Look for the closest enclosing scope that has a reference
      --   -- to `c`.  For now we'll just guess that it is a fully
      --   -- qualified reference.
      --   NodeDef{} -> n
      --   -- Ditto as above.
      --   NodeModule{} -> n
      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 }

findClosest :: ReferenceM m => Set Node -> Node -> m Node
findClosest decls n = do
  Context c <- getContext
  pure $ go $ c
  where
  go :: Node -> Node
  go c = case Set.member (c <> n) decls of
    True  -> c <> n
    False -> case c of
      NodeModule [] -> case n of  -- Last resort
        NodeDef (FQN _ a) -> NodeDef (FQN mempty a)
        x -> x
      _ -> go $ drop c
  drop :: Node -> Node
  drop = \case
    (NodeDef (FQN a _)) -> NodeModule a
    (NodeModule (Namespace [])) -> error "__IMPOSSIBLE__"
    (NodeModule (Namespace xs)) -> NodeModule $ Namespace $ reverse $ Prelude.tail $ reverse xs

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

class References a where
  entries :: ReferenceM 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 . ReferenceM 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.StmtInt 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 (Ruby.Lvasgn _ b) = entries b
-- We don't consider "local variables".  We could do though and say
-- e.g. the local variable is a `Node` with the same "context" as the
-- caller.
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.Int where
  entries = const $ pure ()

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

updateContext :: ReferenceM 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 :: ReferenceM m => Ruby.Atom -> m Node
atomToNode (Ruby.Atom name) = do
  namespace <- getContext >>= pure . \case
    Context (NodeModule c) -> 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 (FQN ns _)) -> ns
  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, args}
    = entries $ Ruby.Def{atom,begin,args}

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
    Ruby.StmtArray{}    -> acc
    -- TODO: Catch-all.  Is this what we want to do?  See e.g. the
    -- note about instance methods in the backlog.
    _                   -> acc

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