module Record
  ( Record (..),
    fromValue,
    toValue,
    lookup,
    Records,
    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 (mapMaybe)
import Data.Text qualified as T
import Exception (DuplicateField (DuplicateField))
import Field qualified as F
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 :: F.Field -> Record J.Value -> Maybe J.Value
lookup (F.Unqualified ks) (Record _ v) =
  lookup' (N.toList ks) v
lookup (F.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]

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

select' :: Records J.Value -> F.Field -> J.Value
select' rs f = case mapMaybe (lookup f) rs of
  [] -> J.Null
  [v] -> v
  (_ : _) -> throw (DuplicateField (F.toString f))

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)