{-# LANGUAGE DeriveFoldable #-}
{-# LANGUAGE DeriveTraversable #-}

-- | Module for a 'FragTree' data structure. This data structure represents
-- "fragments" of a file, which are logical points in a file or buffer.
--
-- A fragment represents an indexed position in a file, but the position
-- is not necessarily discrete. Instead, fragments can be positioned
-- relative to each other (e.g., above or below other fragments),
-- allowing for arbitrary resolution within the file.
module Language.Fiddle.Compiler.Backend.Internal.FragTree
  ( -- * Types
    Fragment,
    FragTree,

    -- * Fragment Constructors
    center,
    -- | A fragment at the center of the file.
    above,
    -- | A fragment above a given fragment.
    below,
    -- | A fragment below a given fragment.
    -- * FragTree Operations
    updateWithReturn,
    -- | Update a fragment in the tree and return a custom value.
    Language.Fiddle.Compiler.Backend.Internal.FragTree.lookup,
    -- | Lookup a value by fragment.
    update,
    -- | Update or insert a value into the tree.
    insertOrReplace,
    -- | Insert a value at a fragment, replacing any existing value.
    insertAppend,
    -- | Insert a value by appending (using 'Monoid') if a value already exists.
    insertWith,
    -- | Insert a value, combining it with an existing value using a custom function.
    delete,
    -- | Delete a value from the tree at a specific fragment.
    singleton,
    -- | Create a tree with a single fragment and value.
  )
where

-- | Direction of a fragment relative to another fragment.
data Loc = Above | Below

-- | A 'Fragment' is a position in a file or buffer that can be relative
-- to other fragments (above or below).
newtype Fragment = Fragment [Loc]

-- | Fragment that represents the center of the file.
center :: Fragment
center = Fragment []

-- | Create a fragment above the given fragment.
above :: Fragment -> Fragment
above (Fragment l) = Fragment $ Above : l

-- | Create a fragment below the given fragment.
below :: Fragment -> Fragment
below (Fragment l) = Fragment $ Below : l

-- | A 'FragTree' represents a binary tree where each node holds a fragment's value.
-- Each tree node has two children, representing the fragments above and below
-- the current fragment. Leaves represent empty nodes.
data FragTree a
  = FragTree (FragTree a) (FragTree a) (Maybe a)
  | FragLeaf
  deriving (Functor)

instance Foldable FragTree where
  foldMap _ FragLeaf = mempty
  foldMap fn (FragTree up down t) =
    foldMap fn up <> foldMap fn t <> foldMap fn down

instance Traversable FragTree where
  traverse _ FragLeaf = pure FragLeaf
  traverse fn (FragTree up down t) = do
    (\a b c -> FragTree a c b)
      <$> traverse fn up
      <*> traverse fn t
      <*> traverse fn down

-- | 'Semigroup' instance for 'FragTree', merging two trees by combining values
-- from corresponding fragments, using the 'Semigroup' instance of the contained values.
instance (Semigroup a) => Semigroup (FragTree a) where
  (<>) (FragTree ab1 be1 v1) (FragTree ab2 be2 v2) =
    FragTree (ab1 <> ab2) (be1 <> be2) (v1 <> v2)
  (<>) FragLeaf f = f
  (<>) f FragLeaf = f

-- | 'Monoid' instance for 'FragTree'. The empty tree is represented by 'FragLeaf'.
instance (Semigroup a) => Monoid (FragTree a) where
  mempty = FragLeaf

-- | 'updateWithReturn' updates a 'FragTree' at the specified 'Fragment',
-- returning both a custom value and the updated tree.
-- If the fragment doesn't exist, it is created.
updateWithReturn :: (Maybe a -> (b, Maybe a)) -> Fragment -> FragTree a -> (b, FragTree a)
updateWithReturn fn (Fragment (reverse -> loc)) = update' fn loc
  where
    -- Helper function for traversing and updating the tree.
    update' fn [] FragLeaf = clean . FragTree FragLeaf FragLeaf <$> fn Nothing
    update' fn [] (FragTree a b v) = clean . FragTree a b <$> fn v
    update' fn (Above : ls) (FragTree a b v) =
      fmap clean $
        (\a' -> FragTree a' b v) <$> update' fn ls a
    update' fn (Below : ls) (FragTree a b v) =
      fmap clean $
        (\b' -> FragTree a b' v) <$> update' fn ls b
    update' fn (Above : ls) FragLeaf =
      fmap clean $
        (\a' -> FragTree a' FragLeaf Nothing) <$> update' fn ls FragLeaf
    update' fn (Below : ls) FragLeaf =
      fmap clean $
        (\b' -> FragTree FragLeaf b' Nothing) <$> update' fn ls FragLeaf

-- | Create a 'FragTree' with a single 'Fragment' and associated value.
singleton :: Fragment -> a -> FragTree a
singleton f a = insertOrReplace f a FragLeaf

-- | Lookup a value in the 'FragTree' by its 'Fragment'. Returns 'Nothing'
-- if the fragment does not exist.
lookup :: Fragment -> FragTree a -> Maybe a
lookup l = fst . updateWithReturn (\a -> (a, a)) l

-- | Update a 'FragTree' at the given 'Fragment'. If the fragment does not exist,
-- it is created. The function is applied to the existing value (or 'Nothing' if no value).
update :: (Maybe a -> Maybe a) -> Fragment -> FragTree a -> FragTree a
update fn fr = snd . updateWithReturn (\ma -> ((), fn ma)) fr

-- | Insert a value into the 'FragTree', replacing any existing value at the given fragment.
insertOrReplace :: Fragment -> a -> FragTree a -> FragTree a
insertOrReplace l a = update (const $ Just a) l

-- | Delete a fragment from the 'FragTree', setting its value to 'Nothing'.
delete :: Fragment -> FragTree a -> FragTree a
delete = update (const Nothing)

-- | Insert a value by appending (using the 'Monoid' instance) to an existing value
-- at the given fragment. If no value exists, the provided value is inserted.
insertAppend :: (Monoid a) => Fragment -> a -> FragTree a -> FragTree a
insertAppend l a = update (<> Just a) l

-- | Insert a value using a custom function to combine it with an existing value
-- at the given fragment. If no val
insertWith :: (a -> a -> a) -> Fragment -> a -> FragTree a -> FragTree a
insertWith fn fr v =
  update
    ( \case
        Nothing -> Just v
        Just a -> Just (fn a v)
    )
    fr

-- | Internal function for cleaning up a 'FragTree' by collapsing empty nodes
-- (i.e., 'FragTree' nodes with no values and no children) into 'FragLeaf'.
clean :: FragTree a -> FragTree a
clean (FragTree FragLeaf FragLeaf Nothing) = FragLeaf
clean x = x