{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DerivingStrategies #-}

module Data.List.NonEmpty.Zipper
  ( Zipper

  -- * Accessors
  , lefts
  , rights
  , current

  -- * Traversal
  , left
  , right
  , findLeft
  , findRight
  , start
  , end

  -- * Construction
  , fromNonEmpty
  , fromNonEmptyEnd

  -- ** Update
  , replace
  , delete
  , push
  , pop
  , shift
  , unshift
  , reverse

  -- * Predicates
  , isStart
  , isEnd
  )
where

import Prelude hiding (reverse)
import qualified Prelude

import Control.Comonad
import Control.DeepSeq (NFData)
import qualified Data.List.NonEmpty as NE
import Data.Maybe (fromMaybe)
import GHC.Generics (Generic)
import Safe (headMay, tailMay)

data Zipper a = Zipper [a] a [a]
  deriving stock (Eq, Show, Functor, Generic)
  deriving anyclass (NFData)

instance Foldable Zipper where
  foldMap f (Zipper ls x rs) = foldMap f (Prelude.reverse ls) `mappend` f x `mappend` foldMap f rs

instance Traversable Zipper where
  traverse f (Zipper ls x rs) =
    Zipper
      <$> (Prelude.reverse <$> traverse f (Prelude.reverse ls))
      <*> f x
      <*> traverse f rs

-- | The list zipper is a basic comonad
--
-- This instance allows us to create a zipper of all possible states of
-- traversing the zipper with 'left' and 'right'.
--
-- >>> duplicate $ fromNonEmpty $ NE.fromList [1, 2, 3]
-- Zipper [] (Zipper [] 1 [2,3]) [Zipper [1] 2 [3],Zipper [2,1] 3 []]
--
instance Comonad Zipper where
  extract = current
  duplicate z =
    let dupWith f r =
          case f r of
            Nothing -> [r]
            Just x -> r:dupWith f x
    in Zipper
      (maybe [] (dupWith left) $ left z)
      z
      (maybe [] (dupWith right) $ right z)

-- | Get the current focus of the @'Zipper'@ cursor
--
-- This is a synonym for 'Control.Comonad.extract'
--
-- >>> current . fromNonEmpty $ NE.fromList [1, 2, 3]
-- 1
--
current :: Zipper a -> a
current (Zipper _ curr _) = curr

-- | Get all values on the left of the cursor
--
-- >>> lefts . fromNonEmptyEnd $ NE.fromList [1, 2, 3]
-- [1,2]
--
lefts :: Zipper a -> [a]
lefts (Zipper ls _ _) = Prelude.reverse ls

-- | Get all values on the right of the cursor
--
-- >>> rights . fromNonEmpty $ NE.fromList [1, 2, 3]
-- [2,3]
--
rights :: Zipper a -> [a]
rights (Zipper _ _ rs) = rs


-- | Move the current focus of the cursor to the left
--
-- >>> left . fromNonEmptyEnd $ NE.fromList [1, 2, 3]
-- Just (Zipper [1] 2 [3])
--
left :: Zipper a -> Maybe (Zipper a)
left (Zipper ps curr ns) = do
  newCurr <- headMay ps
  pure $ Zipper (fromMaybe [] $ tailMay ps) newCurr (curr : ns)

-- | Move the current focus of the cursor to the right
--
-- >>> right . fromNonEmpty $ NE.fromList [1, 2, 3]
-- Just (Zipper [1] 2 [3])
--
right :: Zipper a -> Maybe (Zipper a)
right (Zipper ps curr ns) = do
  newCurr <- headMay ns
  pure $ Zipper (curr : ps) newCurr (fromMaybe [] $ tailMay ns)

-- | Move the current focus of the cursor to the first occurence of a value on the left
--
-- >>> findLeft 2 . fromNonEmptyEnd $ NE.fromList [2, 1, 2, 1, 1, 3]
-- Just (Zipper [1,2] 2 [1,1,3])
--
findLeft :: Eq a => a -> Zipper a -> Maybe (Zipper a)
findLeft target z@(Zipper ps curr ns)
  | curr == target = Just z
  | otherwise = case ps of
    [] -> Nothing
    (x : xs) -> findLeft target (Zipper xs x (curr : ns))

-- | Move the current focus of the cursor to the first occurence of a value on the right
--
-- >>> findRight 3 . fromNonEmpty $ NE.fromList [2, 1, 3, 1, 1, 3]
-- Just (Zipper [1,2] 3 [1,1,3])
--
findRight :: Eq a => a -> Zipper a -> Maybe (Zipper a)
findRight target z@(Zipper ps curr ns)
  | curr == target = Just z
  | otherwise = case ns of
    [] -> Nothing
    (x : xs) -> findRight target (Zipper (curr : ps) x xs)

-- | Move the current focus of the cursor to the start of the @'Zipper'@
--
-- >>> start . fromNonEmptyEnd $ NE.fromList [1, 2, 3]
-- Zipper [] 1 [2,3]
--
start :: Zipper a -> Zipper a
start z
  | isStart z = z
  | otherwise = fromNonEmpty $ toNonEmpty z

-- | Move the current focus of the cursor to the end of the @'Zipper'@
--
-- >>> end . fromNonEmpty $ NE.fromList [1, 2, 3]
-- Zipper [2,1] 3 []
--
end :: Zipper a -> Zipper a
end z
  | isEnd z = z
  | otherwise = fromNonEmptyEnd $ toNonEmpty z


fromNonEmpty :: NE.NonEmpty a -> Zipper a
fromNonEmpty ne = Zipper [] (NE.head ne) (NE.tail ne)

fromNonEmptyEnd :: NE.NonEmpty a -> Zipper a
fromNonEmptyEnd ne = Zipper (NE.tail reversed) (NE.head reversed) []
  where reversed = NE.reverse ne

toNonEmpty :: Zipper a -> NE.NonEmpty a
toNonEmpty (Zipper ls x rs) = NE.fromList $ Prelude.reverse ls ++ [x] ++ rs


-- | Replace the current item under the curosr
--
-- >>> replace 4 . fromNonEmpty $ NE.fromList [1, 2, 3]
-- Zipper [] 4 [2,3]
--
replace :: a -> Zipper a -> Zipper a
replace x (Zipper ls _ rs) = Zipper ls x rs

-- | Delete the item currently under the cursor
--
-- The item currently under the cursor is removed. The cursors focus will move
-- right. If at the end of the @'Zipper'@ the cursor will move left.
--
-- >>> delete . fromNonEmpty $ NE.fromList [1, 2, 3]
-- Just (Zipper [] 2 [3])
--
-- >>> delete . fromNonEmptyEnd $ NE.fromList [1, 2, 3]
-- Just (Zipper [1] 2 [])
--
delete :: Zipper a -> Maybe (Zipper a)
delete (Zipper [] _ []) = Nothing
delete (Zipper ls _ (r : rs)) = Just $ Zipper ls r rs
delete (Zipper (l : ls) _ rs) = Just $ Zipper ls l rs

-- | Insert a value to the left of the cursor
--
-- >>> push 0 . fromNonEmpty $ NE.fromList [1, 2, 3]
-- Zipper [0] 1 [2,3]
--
push :: a -> Zipper a -> Zipper a
push l (Zipper ls x rs) = Zipper (l : ls) x rs

-- | Remove a value to the left of the cursor
--
-- >>> pop . fromNonEmpty $ NE.fromList [1, 2, 3]
-- (Zipper [] 1 [2,3],Nothing)
--
-- >>> pop . fromNonEmptyEnd $ NE.fromList [1, 2, 3]
-- (Zipper [1] 3 [],Just 2)
--
pop :: Zipper a -> (Zipper a, Maybe a)
pop (Zipper [] x rs) = (Zipper [] x rs, Nothing)
pop (Zipper (l : ls) x rs) = (Zipper ls x rs, Just l)

-- | Remove a value to the right of the cursor
--
-- >>> shift . fromNonEmptyEnd $ NE.fromList [1, 2, 3]
-- (Zipper [2,1] 3 [],Nothing)
--
-- >>> shift . fromNonEmpty $ NE.fromList [1, 2, 3]
-- (Zipper [] 1 [3],Just 2)
--
shift :: Zipper a -> (Zipper a, Maybe a)
shift (Zipper ls x []) = (Zipper ls x [], Nothing)
shift (Zipper ls x (r : rs)) = (Zipper ls x rs, Just r)

-- | Insert a value to the right of the cursor
--
-- >>> unshift 4 . fromNonEmpty $ NE.fromList [1, 2, 3]
-- Zipper [] 1 [4,2,3]
--
unshift :: a -> Zipper a -> Zipper a
unshift r (Zipper ls x rs) = Zipper ls x (r : rs)

-- | Reverse the zipper keeping the cursor focus intact
--
-- >>> reverse . fromNonEmpty $ NE.fromList [1, 2, 3]
-- Zipper [2,3] 1 []
--
reverse :: Zipper a -> Zipper a
reverse (Zipper ls x rs) = Zipper rs x ls


-- | Determine if the @'Zipper'@ is at the beginning
--
-- >>> isStart . fromNonEmpty $ NE.fromList [1, 2, 3]
-- True
--
-- >>> isStart . fromNonEmptyEnd $ NE.fromList [1, 2, 3]
-- False
--
isStart :: Zipper a -> Bool
isStart (Zipper [] _ _) = True
isStart _ = False

-- | Determine if the @'Zipper'@ is at the end
--
-- >>> isEnd . fromNonEmptyEnd $ NE.fromList [1, 2, 3]
-- True
--
-- >>> isEnd . fromNonEmpty $ NE.fromList [1, 2, 3]
-- False
--
isEnd :: Zipper a -> Bool
isEnd (Zipper _ _ []) = True
isEnd _ = False