-- |A function to tie together two sorted Haskell Iterators
module Pipes.OrderedZip (orderedZip, orderedZipAll, orderCheckPipe, WrongInputOrderException(..)) where

import           Control.Exception      (Exception)
import           Control.Monad.IO.Class (MonadIO, liftIO)
import           Data.IORef             (newIORef, readIORef, writeIORef)
import           Data.Maybe             (catMaybes)
import           Pipes                  (Pipe, Producer, cat, for, lift, next,
                                         yield, (>->))
import qualified Pipes.Prelude          as P
import           Pipes.Safe             (MonadSafe, throwM)

-- |orderedZip takes a comparison function and two producers and merges them
-- together, creating a new Producer that yields pairs of Maybes of the two
-- datatables provided by the two original producers.
-- The output pairs reflect the Union of the two input producers, with Nothings indicating
-- missing data in one of the producers at that location.
orderedZip :: (Monad m) => (a -> b -> Ordering) -- ^The function to compare types of a with b
           -> Producer a m r1 -- ^The first producer (assumed to be ordered)
           -> Producer b m r2 -- ^The second producer (assumed to be ordered)
           -> Producer (Maybe a, Maybe b) m (r1, r2) -- ^The merged producer
orderedZip ord p1 p2 = do
    p1Front <- lift $ next p1
    p2Front <- lift $ next p2
    go ord p1Front p1 p2Front p2
    where
        go ord' p1Front p1' p2Front p2' = case (p1Front, p2Front) of
            (Left p1r, Left p2r) -> return (p1r, p2r)
            (Left _, Right (p2a, p2Rest)) -> do
                yield (Nothing, Just p2a)
                p2Front' <- lift $ next p2Rest
                go ord' p1Front p1' p2Front' p2Rest
            (Right (p1a, p1Rest), Left _) -> do
                yield (Just p1a, Nothing)
                p1Front' <- lift $ next p1Rest
                go ord' p1Front' p1Rest p2Front p2'
            (Right (p1a, p1Rest), Right (p2a, p2Rest)) -> case ord p1a p2a of
                LT -> do
                    yield (Just p1a, Nothing)
                    p1Front' <- lift $ next p1Rest
                    go ord' p1Front' p1Rest p2Front p2'
                EQ -> do
                    yield (Just p1a, Just p2a)
                    p1Front' <- lift $ next p1Rest
                    p2Front' <- lift $ next p2Rest
                    go ord' p1Front' p1Rest p2Front' p2Rest
                GT -> do
                    yield (Nothing, Just p2a)
                    p2Front' <- lift $ next p2Rest
                    go ord' p1Front p1' p2Front' p2Rest

-- |orderedZipAll takes a comparison function and a list of producers and merges them
-- together, creating a new Producer that yields lists of Maybes of the
-- input data type provided by the original producers.
-- The output list reflects the Union of all input producers, with Nothings indicating
-- missing data in any of the producers at that instance.
orderedZipAll :: (Monad m) => (a -> a -> Ordering) -- ^The function to compare types of a with itself
              -> [Producer a m r] -- ^A list of producers (have to be ordered)
              -> Producer [Maybe a] m [r] -- ^The merged producer
orderedZipAll compFunc prodList = go compFunc (length prodList) prodList
  where
    go :: (Monad m) => (a -> a -> Ordering) -> Int -> [Producer a m r] -> Producer [Maybe a] m [r]
    go compFunc _ [] = error "orderedZipAll error: empty list" -- just to make the function complete
    go compFunc _ [prod] = fmap return (prod >-> P.map (return . Just))
    go compFunc n (prod1:prods) =
        (fmap (\(r, rs) -> (r:rs)) (orderedZip (compFunc2 compFunc) prod1 (go compFunc (n - 1) prods))) >-> P.map mergeMaybeTuples
      where
        mergeMaybeTuples :: (Maybe a, Maybe [Maybe a]) -> [Maybe a]
        mergeMaybeTuples (Nothing, Nothing) = error "orderedZipAll - should never happen" -- just to make this function complete
        mergeMaybeTuples (Nothing, Just mla) = (Nothing: mla)
        mergeMaybeTuples (Just a, Nothing) = (Just a : replicate (n - 1) Nothing)
        mergeMaybeTuples (Just a, Just mla) = (Just a : mla)
        compFunc2 :: (a -> a -> Ordering) -> a -> [Maybe a] -> Ordering
        compFunc2 compFunc x xs =
            let allJusts = catMaybes xs
            in  case allJusts of
                    []      -> error "orderedZipAll compFunc2 - should never happen" -- just to make this complete
                    (xs1:_) -> compFunc x xs1

-- an exception type to represent invalid input order
data WrongInputOrderException = WrongInputOrderException String
    deriving (Show, Eq)
instance Exception WrongInputOrderException

-- a pipe to check wether the stream is ordered according to a custom ordering function
orderCheckPipe :: (MonadIO m, MonadSafe m, Show a) => (a -> a -> Ordering) -- ^the custom ordering function
               -> Pipe a a m r -- ^the resulting pipe
orderCheckPipe cmpFunc = do
    lastValRef <- liftIO $ newIORef (Nothing)
    for cat $ \entry -> do
        lastVal <- liftIO $ readIORef lastValRef
        case lastVal of
            Nothing -> return ()
            Just p -> case cmpFunc entry p of
                LT -> throwM $ WrongInputOrderException ("ordering violated: " ++ show p ++ " should come after " ++ show entry)
                _ -> return ()
        liftIO $ writeIORef lastValRef (Just entry)
        yield entry