module Query.Record
  ( Record (..),
    fromValue,
    toValue,
    lookup,
    Records,
    lookups,
    select,
    disjointUnion,
    disjointUnions,
  )
where

import Control.Exception (throw)
import Data.Aeson qualified as J
import Data.Aeson.Key qualified as JK
import Data.Aeson.KeyMap qualified as JM
import Data.List (foldl')
import Data.List.NonEmpty qualified as N
import Data.Maybe (fromMaybe, mapMaybe)
import Data.Text qualified as T
import Exception (DuplicateField (DuplicateField))
import Query.Field
import Prelude hiding (lookup)

data Record a
  = Record Collection a
  deriving (Show, Eq)

type Collection = FilePath

fromValue :: Collection -> a -> Record a
fromValue = Record

toValue :: Record a -> a
toValue (Record _ v) = v

lookup :: Field -> Record J.Value -> Maybe J.Value
lookup (Unqualified ks) (Record _ v) =
  lookup' (N.toList ks) v
lookup (Qualified c' ks) (Record c v)
  | c' == c = lookup' (N.toList ks) v
  | otherwise = Nothing

lookup' :: [T.Text] -> J.Value -> Maybe J.Value
lookup' [] v = Just v
lookup' (k : ks) (J.Object kvs) =
  lookup' ks =<< JM.lookup (JK.fromText k) kvs

type Records a = [Record a]

lookups :: Field -> Records J.Value -> Maybe J.Value
lookups f rs =
  case mapMaybe (lookup f) rs of
    [] -> Nothing
    [v] -> Just v
    (_ : _) -> throw (DuplicateField (toString f))

select :: [Field] -> Records J.Value -> J.Value
select fs rs =
  foldl'
    union
    (J.Object JM.empty)
    (map (\f -> prefix f ((fromMaybe J.Null (lookups f rs)))) fs)

union :: J.Value -> J.Value -> J.Value
union (J.Object r) (J.Object s) = J.Object (JM.unionWith union r s)

disjointUnion :: J.Value -> J.Value -> J.Value
disjointUnion (J.Object r) (J.Object s) =
  J.Object (JM.unionWithKey disjointUnion' r s)

disjointUnion' :: JK.Key -> J.Value -> J.Value -> J.Value
disjointUnion' _ (J.Object r) (J.Object s) =
  J.Object (JM.unionWithKey disjointUnion' r s)
disjointUnion' k _ _ =
  throw (DuplicateField (JK.toString k))

disjointUnions :: [J.Value] -> J.Value
disjointUnions = foldl' disjointUnion (J.Object JM.empty)