{-# LANGUAGE DeriveFoldable #-}

module Language.Fiddle.Internal.Scopes where

import Control.Monad (forM)
import Data.Aeson
import Data.Aeson.Key
import Data.List (inits, intercalate)
import Data.List.NonEmpty (NonEmpty (..), prependList)
import Data.Map (Map)
import qualified Data.Map as Map
import Data.Maybe (fromMaybe)

-- | 'Scope' represents a hierarchical structure for storing key-value pairs.
--   It can contain nested sub-scopes, which are stored in 'subScopes',
--   and the values for a specific scope are stored in 'scopeValues'.
data Scope k v = Scope
  { subScopes :: Map k (Scope k v), -- Nested sub-scopes
    scopeValues :: Map k v -- Values stored in the current scope
  }
  deriving (Eq, Ord, Show, Read, Functor, Foldable)

-- | 'ScopePath' keeps track of the current scope path as a list of keys,
--   and also includes any additional paths (like imported modules or
--   using namespaces) that might be referenced for lookup.
data ScopePath k = ScopePath
  { currentScope :: [k], -- Current path within the scope hierarchy
    usingPaths :: [[k]] -- Additional paths for resolving symbols
  }
  deriving (Eq, Ord, Show, Read)

-- | The 'Semigroup' instance for 'Scope' allows combining two scopes,
--   where sub-scopes and values are merged together.
instance (Ord k) => Semigroup (Scope k t) where
  (Scope ss1 sv1) <> (Scope ss2 sv2) =
    Scope (Map.unionWith (<>) ss1 ss2) (Map.union sv1 sv2)

-- | The 'Monoid' instance for 'Scope' provides an empty scope with
--   no sub-scopes or values.
instance (Ord k) => Monoid (Scope k t) where
  mempty = Scope mempty mempty

-- | The 'Semigroup' instance for 'ScopePath' allows combining paths,
--   appending the current scope path and using paths.
instance Semigroup (ScopePath k) where
  (ScopePath a1 b1) <> (ScopePath a2 b2) = ScopePath (a1 <> a2) (b1 <> b2)

-- | The 'Monoid' instance for 'ScopePath' provides an empty path.
instance Monoid (ScopePath k) where
  mempty = ScopePath mempty mempty

-- | 'upsertScope' attempts to insert a value 'v' into the 'Scope' at the given
--   key path ('NonEmpty k'). If the key path already exists, the value at the
--   final key is replaced, and the original value is returned in the result.
--   If the key path does not exist, it is created. The function returns a tuple
--   containing the previous value (if any) and the updated scope.
--
--   This function effectively performs an "insert-or-update" operation, allowing
--   you to upsert values into nested scopes while tracking any existing value
--   that was replaced.
upsertScope :: (Ord k) => NonEmpty k -> t -> Scope k t -> (Maybe t, Scope k t)
upsertScope (s :| []) v (Scope ss sv) =
  Scope ss <$> Map.insertLookupWithKey (\_ n _ -> n) s v sv
upsertScope (s :| (a : as)) v (Scope ss sv) =
  let subscope = fromMaybe mempty (Map.lookup s ss)
      (replaced, subscope') = upsertScope (a :| as) v subscope
   in (replaced, Scope (Map.insert s subscope' ss) sv)

insertScope :: (Ord k) => NonEmpty k -> t -> Scope k t -> Scope k t
insertScope p s = snd . upsertScope p s

-- insertScope :: (Ord k) => NonEmpty k -> t -> Scope k t -> Scope k t
-- insertScope a b = snd . upsertScope a b

-- | 'lookupScope' performs a lookup of a value in the scope using a key path
--   ('NonEmpty k'). It traverses through sub-scopes as defined by the path.
lookupScope :: (Ord k) => NonEmpty k -> Scope k t -> Maybe t
lookupScope (s :| []) (Scope _ sv) = Map.lookup s sv
lookupScope (s :| (a : as)) (Scope ss _) = do
  subscope <- Map.lookup s ss
  lookupScope (a :| as) subscope

-- | 'lookupScopeWithPath' searches for a key in the scope by trying all possible
--   paths, including the current scope path and any additional 'using' paths.
--   It returns a list of all valid matches, each paired with the corresponding
--   full key path that resolved to a value. This is useful in cases where symbol
--   resolution might be ambiguous or multiple valid resolutions exist.
--
--   The result is a list of tuples, where the first element of each tuple is the
--   fully-resolved key path (as a 'NonEmpty' list) and the second element is the
--   resolved value.
--
--   This function handles multiple levels of scope resolution, such as those
--   introduced by 'using' directives, and concatenates the results from each
--   possible path.
--
--   Parameters:
--   * 'ScopePath k' - The current scope path and any additional 'using' paths.
--   * 'NonEmpty k'  - The key path to resolve within the scope.
--   * 'Scope k t'   - The scope containing sub-scopes and values.
--
--   Returns: A list of tuples where each tuple contains a fully-resolved key
--   path (as a 'NonEmpty k') and the corresponding value ('t').
lookupScopeWithPath ::
  (Ord k) =>
  ScopePath k ->
  NonEmpty k ->
  Scope k t ->
  [(NonEmpty k, t)]
lookupScopeWithPath (ScopePath current others) key scope =
  let allPaths = reverse (inits current) ++ others
   in flip concatMap allPaths $ \prefix -> do
        case lookupScope (prependList prefix key) scope of
          Just s -> [(prependList prefix key, s)]
          Nothing -> []

instance (ToJSONKey k, ToJSON v, Ord k) => ToJSON (Scope k v) where
  toJSON (Scope subScopes scopeValues) =
    object
      [ fromString "subScopes" .= toJSON subScopes,
        fromString "scopeValues" .= toJSON scopeValues
      ]

instance (FromJSONKey k, FromJSON v, Ord k) => FromJSON (Scope k v) where
  parseJSON (Object v) =
    Scope
      <$> v .: fromString "subScopes"
      <*> v .: fromString "scopeValues"
  parseJSON _ = fail "Expected an object for Scope"