module Data.List.Replace
    ( replace
    , substAllM
    )
    where

import Control.Applicative ((<$>), (<*>))
import Control.Exception (assert)
import Control.Monad.Writer.Lazy (tell, WriterT)
import Control.Monad.Trans (lift)
import Data.List (isPrefixOf, mapAccumR)
import Data.List.Lazy (intercalate_)
import Data.Monoid (Monoid(..))
import Text.Regex.Base.RegexLike (Extract(..))

-- Based on a function by Joseph
-- | Replace a substring with a replacement string throughout a list
replace :: Eq a => [a] -> [a] -> [a] -> [a]
replace []     newSub = intercalate_ newSub . (return <$>)
replace oldSub newSub = _replace
    where
      _replace list@(h:ts)
          | isPrefixOf oldSub list = newSub ++ _replace (drop len list)
          | otherwise              = h : _replace ts
      _replace [] = []
      len = length oldSub

-- | Specifies a range of indexes within a list. Invariant: fst >= 0 && snd >= 0
type IndexRange = (Int, Int)

eSplitAt :: Extract a => Int -> a -> (a, a)
eSplitAt i x = (before i x, after i x)

-- | Gets (before start, (matched, after end))
splitAtIR :: Extract a => IndexRange -> a -> (a, (a, a))
splitAtIR ir = (eSplitAt (snd ir) <$>) . eSplitAt (fst ir)

isBefore :: IndexRange -> IndexRange -> Bool
isBefore (s1, l1) (s2, _) = s1 + l1 <= s2

sortedBy :: (a -> a -> Bool) -> [a] -> Bool
sortedBy _ [] = True
sortedBy r (h:t) = impl h t
  where
    impl _ [] = True
    impl h' (h2:t2) = r h' h2 && impl h2 t2

-- | Gets all the unmatched and matched segments, in order
-- Precondition: The IndexRange list must be in ascending order
unweave :: Extract a => [IndexRange] -> a -> (a, [(a, a)])
unweave = assert . sortedBy isBefore <*> flip (mapAccumR $ flip splitAtIR)

liftAndTell :: Monad m => m a -> WriterT [a] m a
liftAndTell m = do
  y <- lift m
  tell [y]
  return y

-- | Substitutes all ranges using the values returned by the function provided (which are also returned).
-- Precondition: The IndexRange list must be in ascending order
substAllM :: (Functor m, Monad m, Extract a, Monoid a)
             => (a -> m a) -> a -> [IndexRange] -> WriterT [a] m a
substAllM f l irs = mconcat . (beforeFirst :) <$> mapM process matches
    where
      (beforeFirst, matches) = unweave irs l
      process (matched, unmatched) = (`mappend` unmatched) <$> liftAndTell (f matched)