-- Copyright 2019 Google LLC
--
-- Use of this source code is governed by a BSD-style
-- license that can be found in the LICENSE file or at
-- https://developers.google.com/open-source/licenses/bsd

{-# LANGUAGE MagicHash #-}

module Util where

import Prelude
import qualified Data.Set as Set
import qualified Data.Map.Strict as M
import Control.Applicative
import Control.Monad.Reader
import Control.Monad.State.Strict
import System.CPUTime
import GHC.Base (getTag)
import GHC.Exts ((==#), tagToEnum#)
import Crypto.Hash
import Data.Functor.Identity (Identity(..))
import Data.Maybe (catMaybes, mapMaybe)
import Data.List (sort)
import Data.Hashable (Hashable)
import Data.Store (Store)
import qualified Data.Text as T
import qualified Data.Text.Encoding as T
import qualified Data.List.NonEmpty as NE
import qualified Data.ByteString    as BS
import Data.Foldable
import Data.Text.Prettyprint.Doc (Pretty (..), pretty)
import Data.List.NonEmpty (NonEmpty (..))
import GHC.Generics (Generic)

import Err

class IsBool a where
  toBool :: a -> Bool

iota :: (Enum a, Integral a) => a -> [a]
iota n = take (fromEnum n) [0..] -- XXX: `[0..(n-1)]` is incorrect for unsigned ints!
{-# INLINE iota #-}

swapAt :: Int -> a -> [a] -> [a]
swapAt _ _ [] = error "swapping to empty list"
swapAt 0 y (_:xs) = y:xs
swapAt n y (x:xs) = x:(swapAt (n-1) y xs)

onFst :: (a -> b) -> (a, c) -> (b, c)
onFst f (x, y) = (f x, y)

onFstM :: (Functor m) => (a -> m b) -> (a, c) -> m (b, c)
onFstM f (x, y) = (,y) <$> f x

onSnd :: (a -> b) -> (c, a) -> (c, b)
onSnd f (x, y) = (x, f y)

onSndM :: (Functor m) => (a -> m b) -> (c, a) -> m (c, b)
onSndM f (x, y) = (x,) <$> f y

unsnocNonempty :: NonEmpty a -> ([a], a)
unsnocNonempty (x:|xs) = case reverse (x:xs) of
  (y:ys) -> (reverse ys, y)
  _ -> error "impossible"

enumerate :: Traversable f => f a -> f (Int, a)
enumerate xs = evalState (traverse addCount xs) 0
  where addCount :: a -> State Int (Int, a)
        addCount x = do n <- get
                        put (n + 1)
                        return (n, x)

splitMap :: Ord k => [k] -> M.Map k v -> (M.Map k v, M.Map k v)
splitMap ks m = let ks' = Set.fromList ks
                    pos = M.filterWithKey (\k _ -> k `Set.member` ks') m
                in (pos, m M.\\ pos)

listDiff :: Ord a => [a] -> [a] -> [a]
listDiff xs ys = Set.toList $ Set.difference (Set.fromList xs) (Set.fromList ys)

showErr :: Show e => Either e a -> Either String a
showErr (Left e) = Left (show e)
showErr (Right x) = Right x

group :: (Ord a) => [(a,b)] -> [(a, [b])]
group [] = []
group ((k,v):xs) =
  let (prefix, suffix) = span ((== k) . fst) xs
      g = v:(map snd prefix)
  in (k, g) : group suffix

ungroup ::  [(a, [b])] -> [(a,b)]
ungroup [] = []
ungroup ((k,vs):xs) = (zip (repeat k) vs) ++ ungroup xs

uncons :: [a] -> (a, [a])
uncons (x:xs) = (x, xs)
uncons [] = error "whoops! [uncons]"

pad :: a -> Int -> [a] -> [a]
pad v n xs = xs ++ replicate (n - length(xs)) v

padLeft :: a -> Int -> [a] -> [a]
padLeft v n xs = replicate (n - length(xs)) v ++ xs

-- Nothing if the exact prefix isn't available
splitAtExact :: Int -> [a] -> Maybe ([a], [a])
splitAtExact n xs =
  if n <= length xs
    then Just $ splitAt n xs
    else Nothing

delIdx :: Int -> [a] -> [a]
delIdx i xs = case splitAt i xs of
  (prefix, _:suffix) -> prefix ++ suffix
  (prefix, []) -> prefix -- Already not there

replaceIdx :: Int -> a -> [a] -> [a]
replaceIdx i x xs = case splitAt i xs of
  (prefix, _:suffix) -> prefix ++ (x:suffix)
  (prefix, []) -> prefix ++ [x]

insertIdx :: Int -> a -> [a] -> [a]
insertIdx i x xs = case splitAt i xs of
  (prefix, suffix) -> prefix ++ (x:suffix)

mvIdx :: Int -> Int -> [a] -> [a]
mvIdx i j xs | j == i = xs
             | j < i = let x = xs!!i
                       in insertIdx j x . delIdx i $ xs
             | otherwise = let x = xs!!i
                           in  delIdx i . insertIdx j x $ xs

mapFst :: (a -> b) -> [(a, c)] -> [(b, c)]
mapFst f zs = [(f x, y) | (x, y) <- zs]

mapSnd :: (a -> b) -> [(c, a)] -> [(c, b)]
mapSnd f zs = [(x, f y) | (x, y) <- zs]

foldJusts :: Monoid b => [a] -> (a -> Maybe b) -> b
foldJusts xs f = fold $ mapMaybe f xs

forMFilter :: Monad m => [a] -> (a -> m (Maybe b)) -> m [b]
forMFilter xs f = catMaybes <$> mapM f xs
{-# INLINE forMFilter #-}

composeN :: Int -> (a -> a) -> a -> a
composeN n f = foldr (.) id (replicate n f)

repeated :: Ord a => [a] -> [a]
repeated = repeatedSorted . sort

repeatedSorted :: Eq a => [a] -> [a]
repeatedSorted [] = []
repeatedSorted [_] = []
repeatedSorted (x:y:rest) | x == y = [x] ++ repeatedSorted (dropWhile (== x) rest)
                          | otherwise = repeatedSorted (y:rest)

splitOn :: (a -> Bool) -> [a] -> [[a]]
splitOn f s = let (prefix, suffix) = break f s
              in case suffix of
                   [] -> [prefix]
                   _:xs -> prefix : splitOn f xs

restructure :: Traversable f => [a] -> f b -> f a
restructure xs structure = evalState (traverse procLeaf structure) xs
  where procLeaf :: b -> State [a] a
        procLeaf _ = do ~(x:rest) <- get
                        put rest
                        return x

-- TODO: find a more efficient implementation
findReplace :: Eq a => [a] -> [a] -> [a] -> [a]
findReplace _ _ [] = []
findReplace old new s@(x:xs) =
  if take n s == old
    then new ++ recur (drop n s)
    else x : recur xs
  where n = length old
        recur = findReplace old new

scan :: Traversable t => (a -> s -> (b, s)) -> t a -> s -> (t b, s)
scan f xs s = runState (traverse (asState . f) xs) s

scanM :: (Monad m, Traversable t) => (a -> s -> m (b, s)) -> t a -> s -> m (t b, s)
scanM f xs s = runStateT (traverse (asStateT . f) xs) s

asStateT :: Monad m => (s -> m (a, s)) -> StateT s m a
asStateT f = do
  s <- get
  (ans, s') <- lift $ f s
  put s'
  return ans

asState :: (s -> (a, s)) -> State s a
asState f = asStateT (Identity . f)

uncurry3 :: (a -> b -> c -> d) -> (a, b, c) -> d
uncurry3 f (x, y, z) = f x y z

bindM2 :: Monad m => (a -> b -> m c) -> m a -> m b -> m c
bindM2 f ma mb = do
  a <- ma
  b <- mb
  f a b

(...) :: (c -> d) -> (a -> b -> c) -> a -> b -> d
f ... g = \x y -> f $ g x y

foldMapM :: (Monad m, Monoid w, Foldable t) => (a -> m w) -> t a -> m w
foldMapM f xs = foldM (\acc x -> (acc<>) <$> f x ) mempty xs

lookupWithIdx :: Eq a => a -> [(a, b)] -> Maybe (Int, b)
lookupWithIdx k vals = lookup k $ [(x, (i, y)) | (i, (x, y)) <- zip [0..] vals]

-- NOTE: (toList args) has to be at least as long as (toList trav)
zipWithT :: (Traversable t, Monad h, Foldable f) => (a -> b -> h c) -> t a -> f b -> h (t c)
zipWithT f trav args = flip evalStateT (toList args) $ flip traverse trav $ \e -> getNext >>= lift . f e
  where getNext = get >>= \(h:t) -> put t >> return h

for :: Functor f => f a -> (a -> b) -> f b
for = flip fmap

transitiveClosure :: forall a. Ord a => (a -> [a]) -> [a] -> [a]
transitiveClosure getParents seeds =
  toList $ execState (mapM_ go seeds) mempty
  where
    go :: a -> State (Set.Set a) ()
    go x = do
      visited <- get
      unless (x `Set.member` visited) $ do
        modify $ Set.insert x
        mapM_ go $ getParents x

transitiveClosureM :: forall m a. (Monad m, Ord a) => (a -> m [a]) -> [a] -> m [a]
transitiveClosureM getParents seeds =
  toList <$> execStateT (mapM_ go seeds) mempty
  where
    go :: a -> StateT (Set.Set a) m ()
    go x = do
      visited <- get
      unless (x `Set.member` visited) $ do
        modify (<> Set.singleton x)
        lift (getParents x) >>= mapM_ go

measureSeconds :: MonadIO m => m a -> m (a, Double)
measureSeconds m = do
  t1 <- liftIO $ getCPUTime
  ans <- m
  t2 <- liftIO $ getCPUTime
  return (ans, (fromIntegral $ t2 - t1) / 1e12)

whenM :: Monad m => m Bool -> m () -> m ()
whenM test doit = test >>= \case
  True -> doit
  False -> return ()

anyM :: Monad m => (a -> m Bool) -> [a] -> m Bool
anyM f xs = do
  conds <- mapM f xs
  return $ any id conds

allM :: Monad m => (a -> m Bool) -> [a] -> m Bool
allM f xs = do
  conds <- mapM f xs
  return $ all id conds

fromMaybeM :: Monad m => Maybe a -> b -> (a -> m b) -> m b
fromMaybeM optVal defaultVal cont = case optVal of
  Just x -> cont x
  Nothing -> return defaultVal
{-# INLINE fromMaybeM #-}

liftMaybe :: Alternative m => Maybe a -> m a
liftMaybe Nothing = empty
liftMaybe (Just x) = pure x
{-# INLINE liftMaybe #-}

-- === zippable class ===

class Zippable f where
  zipWithZ :: MonadFail m => (a -> b -> m c) -> f a -> f b -> m (f c)

instance Zippable [] where
  zipWithZ _ [] [] = return []
  zipWithZ f (x:xs) (y:ys) = (:) <$> f x y <*> zipWithZ f xs ys
  zipWithZ _ _ _ = zipErr

instance Zippable NE.NonEmpty where
  zipWithZ f xs ys = NE.fromList <$> zipWithZ f (NE.toList xs) (NE.toList ys)

zipWithZ_ :: Zippable f => MonadFail m => (a -> b -> m c) -> f a -> f b -> m ()
zipWithZ_ f xs ys = zipWithZ f xs ys >> return ()

zipErr :: MonadFail m => m a
zipErr = fail $ "zip error. Call stack:\n" ++ printCurrentCallStack (getCurrentCallStack ())

forMZipped :: Zippable f => MonadFail m => f a -> f b -> (a -> b -> m c) -> m (f c)
forMZipped xs ys f = zipWithZ f xs ys

forMZipped_ :: Zippable f => MonadFail m => f a -> f b -> (a -> b -> m c) -> m ()
forMZipped_ xs ys f = void $ forMZipped xs ys f

getAlternative :: Alternative m => [a] -> m a
getAlternative xs = asum $ map pure xs
{-# INLINE getAlternative #-}

newtype SnocList a = ReversedList { fromReversedList :: [a] }
        deriving (Show, Eq, Ord, Generic, Functor) -- XXX: NOT deriving order-sensitive things like Monoid, Applicative etc

instance Semigroup (SnocList a) where
  (ReversedList x) <> (ReversedList y) = ReversedList $ y ++ x
  {-# INLINE (<>) #-}

instance Monoid (SnocList a) where
  mempty = ReversedList []
  {-# INLINE mempty #-}

instance Foldable SnocList where
  foldMap f (ReversedList xs) = foldMap f (reverse xs)
  {-# INLINE foldMap #-}

instance Traversable SnocList where
  traverse f (ReversedList xs) = ReversedList . reverse <$> traverse f (reverse xs)
  {-# INLINE traverse #-}

snoc :: SnocList a -> a -> SnocList a
snoc (ReversedList xs) x = ReversedList (x:xs)
{-# INLINE snoc #-}

emptySnocList :: SnocList a
emptySnocList = ReversedList []
{-# INLINE emptySnocList #-}

unsnoc :: SnocList a -> [a]
unsnoc (ReversedList x) = reverse x
{-# INLINE unsnoc #-}

toSnocList :: [a] -> SnocList a
toSnocList xs = ReversedList $ reverse xs
{-# INLINE toSnocList #-}

tryUnsnoc :: SnocList a -> Maybe (SnocList a, a)
tryUnsnoc (ReversedList []) = Nothing
tryUnsnoc (ReversedList (x:xs)) = Just (ReversedList xs, x)
{-# INLINE tryUnsnoc #-}

-- === generic tree ===

data Tree a = Leaf a | Branch [Tree a]
     deriving (Show, Eq, Ord, Generic, Functor, Foldable, Traversable)
instance Store a => Store (Tree a)
instance Hashable a => Hashable (Tree a)

zipTrees :: Tree a -> Tree b -> Tree (a, b)
zipTrees (Leaf x) (Leaf y) = Leaf (x, y)
zipTrees (Branch xs) (Branch ys) | length xs == length ys = Branch $ zipWith zipTrees xs ys
zipTrees _ _ = error "zip error"

instance Pretty a => Pretty (Tree a) where
  pretty = \case
    Leaf x -> pretty x
    Branch xs -> pretty xs

readFileText :: MonadIO m => FilePath -> m T.Text
readFileText fname = liftIO $ T.decodeUtf8 <$> BS.readFile fname

-- === bytestrings paired with their hash digest ===

-- TODO: use something other than a string to store the digest
type FileHash     = String
type FileContents = BS.ByteString

-- TODO: consider adding mtime as well for a fast path that doesn't
-- require reading the file
data File = File
  { fContents :: FileContents
  , fHash     :: FileHash }
  deriving (Show, Eq, Ord)

addHash :: FileContents -> File
addHash s = File s $ show (hash s :: Digest SHA256)

readFileWithHash :: MonadIO m => FilePath -> m File
readFileWithHash path = liftIO $ addHash <$> BS.readFile path

sameConstructor :: a -> a -> Bool
sameConstructor x y = tagToEnum# (getTag x ==# getTag y)
{-# INLINE sameConstructor #-}