{-# LANGUAGE DeriveFoldable #-}

module Language.Fiddle.Internal.Scopes where

import Data.Foldable
import Data.Aeson
import Data.Aeson.Key
import Data.List (inits)
import Data.List.NonEmpty (NonEmpty (..), prependList)
import Data.Map (Map)
import qualified Data.Map as Map
import Data.Maybe (fromMaybe, maybeToList)
import Control.Arrow (Arrow(first))

-- | '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)

-- | Qualify a name with the current scope.
qualifyPath :: ScopePath k -> NonEmpty k -> NonEmpty k
qualifyPath ScopePath {currentScope = scope} = prependList scope

-- | Push a new scope onto the current scope.
pushScope :: NonEmpty k -> ScopePath k -> ScopePath k
pushScope v s@ScopePath {currentScope = scope} =
  s {currentScope = scope ++ toList v}

-- | Adds a path to the "using" paths.
addUsingPath :: NonEmpty k -> ScopePath k -> ScopePath k
addUsingPath path s@ScopePath {usingPaths = paths} =
  s {usingPaths = toList path : paths}

-- | 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.unionWith (<>) 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 -> ([t], Scope k t)
upsertScope (s :| []) v (Scope ss sv) =
  first (fromMaybe []) $
    Scope ss <$> Map.insertLookupWithKey (const (<>)) 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

singleton :: (Ord k) => NonEmpty k -> t -> Scope k t
singleton ks t = insertScope ks t mempty

-- 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 -> [t]
lookupScope (s :| []) (Scope _ sv) = fromMaybe [] $ Map.lookup s sv
lookupScope (s :| (a : as)) (Scope ss _) = do
  subscope <- maybeToList $ 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 do
        prefix <- allPaths
        let qualifiedKey = prependList prefix key
        value <- lookupScope qualifiedKey scope
        return (qualifiedKey, value)

-- 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 =
    object
      [ fromString "subScopes" .= toJSON (subScopes scope),
        fromString "scopeValues" .= toJSON (scopeValues scope)
      ]

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"