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

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 :: Block -> 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

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 = coerce
  toList = coerce
deriving newtype instance Semigroup Namespace
deriving newtype instance Monoid Namespace
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 = NodeFunction FQN | NodeModule Namespace

deriving stock instance Show Node
deriving stock instance Eq Node
deriving stock instance Ord Node

instance ToJSON Node where
  toJSON = \case
    NodeFunction q -> Aeson.toJSON q
    NodeModule m -> Aeson.toJSON m

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

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

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

newtype Result = Result (Map Node (Set Node))
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
  NodeFunction 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 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 (Context (NodeModule q)) = Context $ NodeModule $ name2ns n <> q
  go (Context NodeFunction{}) = 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 $ Map.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 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
    StmtConst c    -> entries c
    StmtRBlock b   -> entries b
    StmtCasgn c    -> entries c
    StmtArray a    -> entries a
    StmtAnything a -> entries a

-- TODO
instance References Ruby.RBlock where
  entries = const $ pure ()
instance References Ruby.Casgn where
  entries = const $ pure ()
instance References Ruby.RArray where
  entries = const $ pure ()
instance References Ruby.Anything where
  entries = const $ pure ()

instance References Ruby.Namespace where
  entries = \case
    Ruby.Namespace xs -> do
      Context c <- getContext
      let
        ctxt = case c of
          NodeFunction FQN{namespace} -> namespace
          NodeModule namespace -> namespace
        -- TODO Hacky:
        ns = Namespace ((\(Name (Aeson.String t)) -> t) <$> xs)
      -- TODO: Broken
      -- application c (NodeModule $ ctxt `onTop` ns)
      application c (NodeModule $ ns)

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

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


onTop' :: Eq a => [a] -> [a] -> [a]
onTop' [] y = y
onTop' x [] = x
onTop' (x:xss) ys@(y:yss) = if
  | x == y    -> pure x <> yss
  | otherwise -> pure x <> xss `onTop'` ys

onTop :: Namespace -> Namespace -> Namespace
onTop (Namespace xs) (Namespace ys) = Namespace $ reverse $ reverse xs `onTop'` reverse ys

instance References Send where
  entries Send{namespace, name} = do
    Context c <- getContext
    let
      ctxt = case c of
        NodeFunction FQN{namespace = ns} -> ns
        NodeModule ns -> ns
    -- TODO: Broken
    -- application c $ NodeFunction $ qual (ctxt `onTop` fromNS namespace) name
    application c $ NodeFunction $ 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