-- |
--
-- Module      : Data.Aeson.Optics.Ext
-- Copyright   : (c) 2025 Patrick Brisbin
-- License     : AGPL-3
-- Maintainer  : pbrisbin@gmail.com
-- Stability   : experimental
-- Portability : POSIX
module Data.Aeson.Optics.Ext
  ( atKey
  , atNth
  ) where

import Prelude

import Data.Aeson (Key, Value (..))
import Data.Aeson.Key qualified as Key
import Data.Aeson.KeyMap qualified as KeyMap
import Data.Aeson.Optics
import Data.Vector qualified as V
import Data.Vector.Ext qualified as V
import Optics

-- | Like 'key', but uses 'at' instead of 'ix'. This is handy when adding and
-- removing object keys:
--
-- >>> "{\"a\": 100, \"b\": 200}" & atKey "a" .~ Nothing
-- "{\"b\":200}"
--
-- >>> "{\"a\": 100, \"b\": 200}" & atKey "c" ?~ String "300"
-- "{\"a\":100,\"b\":200,\"c\":\"300\"}"
atKey :: Key -> AffineTraversal' Value (Maybe Value)
atKey :: Key -> AffineTraversal' Value (Maybe Value)
atKey Key
k = Prism' Value (KeyMap Value)
forall t. AsValue t => Prism' t (KeyMap Value)
_Object Prism' Value (KeyMap Value)
-> Optic
     A_Lens
     NoIx
     (KeyMap Value)
     (KeyMap Value)
     (Maybe Value)
     (Maybe Value)
-> AffineTraversal' Value (Maybe Value)
forall k l m (is :: IxList) (js :: IxList) (ks :: IxList) s t u v a
       b.
(JoinKinds k l m, AppendIndices is js ks) =>
Optic k is s t u v -> Optic l js u v a b -> Optic m ks s t a b
% Index (KeyMap Value)
-> Lens' (KeyMap Value) (Maybe (IxValue (KeyMap Value)))
forall m. At m => Index m -> Lens' m (Maybe (IxValue m))
at Key
Index (KeyMap Value)
k

-- | Like 'atKey', but for 'Array's
--
-- Adding shifts all later elements right:
--
-- >>> ['a', 'b'] & atNth 1 ?~ 'x'
-- ['a', 'x', 'b']
--
-- Removing shifts all later elements left:
--
-- >>> ['a', 'b', 'c'] & atNth 1 .~ Nothing
-- ['a', 'c']
--
-- __NOTE__: this function will also index objects, in which case this behaves
-- exactly like 'atKey'. This is necessary for our use-case and probably means
-- we could never upstream this.
--
-- >>> {"0": 'a', "1": 'b'} & atNth 1 ?~ 'x'
-- {"0": 'a', "1": 'x', "2": 'b'}
atNth :: Int -> AffineTraversal' Value (Maybe Value)
atNth :: Int -> AffineTraversal' Value (Maybe Value)
atNth Int
n = (Value -> Either Value (Maybe Value))
-> (Value -> Maybe Value -> Value)
-> AffineTraversal' Value (Maybe Value)
forall s t a b.
(s -> Either t a) -> (s -> b -> t) -> AffineTraversal s t a b
atraversal Value -> Either Value (Maybe Value)
matcher Value -> Maybe Value -> Value
updater
 where
  matcher :: Value -> Either Value (Maybe Value)
  matcher :: Value -> Either Value (Maybe Value)
matcher = \case
    Object KeyMap Value
km -> Maybe Value -> Either Value (Maybe Value)
forall a b. b -> Either a b
Right (Maybe Value -> Either Value (Maybe Value))
-> Maybe Value -> Either Value (Maybe Value)
forall a b. (a -> b) -> a -> b
$ Key -> KeyMap Value -> Maybe Value
forall v. Key -> KeyMap v -> Maybe v
KeyMap.lookup (String -> Key
Key.fromString (String -> Key) -> String -> Key
forall a b. (a -> b) -> a -> b
$ Int -> String
forall a. Show a => a -> String
show Int
n) KeyMap Value
km
    Array Array
vec -> Maybe Value -> Either Value (Maybe Value)
forall a b. b -> Either a b
Right (Maybe Value -> Either Value (Maybe Value))
-> Maybe Value -> Either Value (Maybe Value)
forall a b. (a -> b) -> a -> b
$ Array
vec Array -> Int -> Maybe Value
forall a. Vector a -> Int -> Maybe a
V.!? Int
n
    Value
v -> Value -> Either Value (Maybe Value)
forall a b. a -> Either a b
Left Value
v

  updater :: Value -> Maybe Value -> Value
  updater :: Value -> Maybe Value -> Value
updater Value
nv = \case
    Maybe Value
Nothing -> case Value
nv of
      Object KeyMap Value
km -> KeyMap Value -> Value
Object (KeyMap Value -> Value) -> KeyMap Value -> Value
forall a b. (a -> b) -> a -> b
$ Key -> KeyMap Value -> KeyMap Value
forall v. Key -> KeyMap v -> KeyMap v
KeyMap.delete (String -> Key
Key.fromString (String -> Key) -> String -> Key
forall a b. (a -> b) -> a -> b
$ Int -> String
forall a. Show a => a -> String
show Int
n) KeyMap Value
km
      Array Array
vec -> Array -> Value
Array (Array -> Value) -> Array -> Value
forall a b. (a -> b) -> a -> b
$ Int -> Array -> Array
forall a. Int -> Vector a -> Vector a
V.deleteAt Int
n Array
vec
      Value
v -> Value
v
    Just Value
x -> case Value
nv of
      Object KeyMap Value
km -> KeyMap Value -> Value
Object (KeyMap Value -> Value) -> KeyMap Value -> Value
forall a b. (a -> b) -> a -> b
$ Key -> Value -> KeyMap Value -> KeyMap Value
forall v. Key -> v -> KeyMap v -> KeyMap v
KeyMap.insert (String -> Key
Key.fromString (String -> Key) -> String -> Key
forall a b. (a -> b) -> a -> b
$ Int -> String
forall a. Show a => a -> String
show Int
n) Value
x KeyMap Value
km
      Array Array
vec -> Array -> Value
Array (Array -> Value) -> Array -> Value
forall a b. (a -> b) -> a -> b
$ Int -> Value -> Array -> Array
forall a. Int -> a -> Vector a -> Vector a
V.insertAt Int
n Value
x Array
vec
      Value
v -> Value
v