{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}

-- | Read and write a record into JSON tree
module Data.JSOP where

import Control.Lens (APrism', clonePrism, re, (%~), (.~), (^.), (^?), _1)
import Data.Aeson (Value (..), object)
import Data.Aeson.Lens
import qualified Data.HashMap.Strict as HM
import qualified Data.Map as M
import qualified Data.Map.Monoidal.Strict as Mm
import Generics.SOP
import Protolude hiding (All)
import Trie

-- | prism for a 'Value'
type ValuePrism a = APrism' Value a

-- | A record field operation specification
data JSOP p a = JSOP
  { -- | custom path encoding
    jsop_path :: p
  , -- | a default value , Nothing means value is required
    jsop_default :: Maybe a
  , -- | the prism to read and write the record field
    jsop_prism :: ValuePrism a
  }

-- | shortcut for a parser that handle missing values
required
  :: p -- ^ path
  -> ValuePrism a -- ^ prism
  -> JSOP p a
required path = JSOP path Nothing

-- | shortcut for parsers that have a default if value is missing
optional
  :: p -- ^ path
  -> a -- ^ default value
  -> ValuePrism a -- ^ prism
  -> JSOP p a
optional path = JSOP path . Just

-- | parsing problems
data JSOPIssue
  = -- | the index of the field in the record that failed, with the Value if reacheable and the requested type
    JSOPIssue (Int, Maybe Value, TypeRep)
  | -- | paths are not enough or too many
    JSOPWrongNumberOfPaths
  deriving (Eq, Show)

-- | get a record out of the json tree
jread
  :: (All Typeable xs, IsProductType a xs)
  => (path -> [Text]) -- ^ how to extract keys from a path
  -> NP (JSOP path) xs -- ^ prisms for the path-indexed values
  -> Value -- ^ json structure holding the record
  -> Either JSOPIssue a -- ^ the record , if possible
jread splitter ps value = maybe
  do Left JSOPWrongNumberOfPaths
  do
    fmap productTypeTo
      . hsequence
      . hcliftA2 (Proxy :: Proxy Typeable) parseSField ps
  do fromList $ zip [0 ..] $ getValues splitter (paths ps) value

paths :: All Top xs => NP (JSOP p) xs -> [p]
paths = hcollapse . hmap (K . jsop_path)

-- parse a single 'Value'
parseSField :: forall a p. Typeable a => JSOP p a -> K (Int, Maybe Value) a -> Either JSOPIssue a
parseSField (JSOP _ md parser) (K (n, v)) = case v of
  Nothing -> case md of
    Just x -> Right x
    Nothing -> Left $ JSOPIssue (n, v, typeOf @a $ panic "no value")
  Just w -> case w ^? clonePrism parser of
    Nothing -> Left $ JSOPIssue (n, v, typeOf @a $ panic "cannot parse")
    Just r -> Right r

--------------------------------------------------------------------------------------------
-- memoize a trie to resolve path queries
--------------------------------------------------------------------------------------------
type QPaths = Trie Text (First Int, [Int])

mkPath :: Int -> [Text] -> QPaths
mkPath n =
  foldr
    (fmap (Trie (First Nothing, [n])) . Mm.singleton)
    (Trie (First (Just n), [n]) mempty)

mkPaths :: [[Text]] -> QPaths
mkPaths = foldMap (uncurry mkPath) . zip [0 ..]

treequery :: QPaths -> Value -> ([(Int, Value)], [Int])
treequery (Trie (m, _ns) qs) v =
  appEndo (foldMap (\n -> Endo $ _1 %~ (:) (n, v)) m) $
    fold $ do
      (k, t) <- Mm.assocs qs
      pure $ case v ^? key k of
        Nothing -> ([], snd $ load t)
        Just w -> treequery t w

-- | you should close over paths argument to get an efficient Value -> [Maybe Value]
getValues
  :: (p -> [Text]) -- ^ how to extract keys from a path
  -> [p] -- ^ paths
  -> Value -- ^ json tree
  -> [Maybe Value] -- ^ found at path values
getValues splitter ts v =
  let (positive, negative) = treequery (mkPaths $ splitter <$> ts) v
   in fmap snd $
        sortOn fst $
          fmap (,Nothing) negative <> fmap (fmap Just) positive

type TPaths = Trie Text (First Int, Any)

mkTPath :: Int -> (Bool, [Text]) -> TPaths
mkTPath n (b, xs) =
  foldr
    (fmap (Trie mempty) . Mm.singleton)
    (Trie (First $ Just n, Any b) mempty)
    xs

mkTPaths :: [(Bool, [Text])] -> TPaths
mkTPaths = foldMap (uncurry mkTPath) . zip [0 ..]

treechange :: Bool -> TPaths -> Value -> Map Int Value -> Value
treechange c (Trie (m, Any b) qs) v cs = case m of
  First Nothing -> foldl' (changenode c) v $ Mm.assocs qs
  First (Just n) ->
    if c
      then if b then cs M.! n else Null
      else cs M.! n
  where
    changenode :: Bool -> Value -> (Text, TPaths) -> Value
    changenode c' w (k, rest) = case w ^? key k of
      Nothing -> _Object %~ HM.insert k (treechange True rest (object []) cs) $ w
      Just w' -> key k .~ treechange c' rest w' cs $ w

-- | you should close over paths argument to get an efficient Value -> Value
setValues
  :: (p -> (Bool, [Text])) -- ^ how to extract keys from paths, and if the path is required
  -> [(p, Value)] -- ^ what to substitute at each path
  -> Value -- ^ value to amend
  -> Value -- ^ amended value
setValues splitter ts v =
  treechange False (mkTPaths $ splitter . fst <$> ts) v $
    M.fromList $ zip [0 ..] $ snd <$> ts

-- | not very well defined write at path operation
jwrite
  :: (All Typeable xs, IsProductType a xs, All Top xs, Show path, Ord path)
  => (path -> [Text]) -- ^ how to extract keys from a path
  -> NP (JSOP path) xs -- ^ prisms for the path-indexed values
  -> Value -- ^ json 'Value' to amend
  -> a -- ^ record to write
  -> Value -- ^ amended 'Value'
jwrite splitter ps value x = setValues splitter' setters value
  where
    splitter' k = (optionals M.! k, splitter k)
    optionals = M.fromList $ (\(p, b, _) -> (p, b)) <$> ops
    ops = f ps x
    setters = (\(p, _, v) -> (p, v)) <$> ops
    f :: IsProductType a xs => NP (JSOP path) xs -> a -> [(path, Bool, Value)]
    f ps' x' = hcollapse $ hzipWith g ps' $ productTypeFrom x'
    g :: JSOP path a1 -> I a1 -> K (path, Bool, Value) a1
    g (JSOP p md parser) (I x') = K (p, isJust md, x' ^. re (clonePrism parser))