{-
  Copyright (c) Meta Platforms, Inc. and affiliates.
  All rights reserved.

  This source code is licensed under the BSD-style license found in the
  LICENSE file in the root directory of this source tree.
-}

module Glean.Database.Data
  ( storeSchema
  , retrieveSchema
  , storeUnits
  , retrieveUnits
  , storeSlices
  , retrieveSlices
  ) where

import Data.Binary
import Data.ByteString (ByteString)
import Data.ByteString.Lazy (toStrict, fromStrict)

import Thrift.Protocol.Compact

import Glean.Database.Exception
import Glean.Database.Storage (Storage, Database)
import qualified Glean.Database.Storage as Storage
import Glean.RTS.Foreign.Ownership
import Glean.Types (Repo)
import Glean.Internal.Types (StoredSchema)

sCHEMA_KEY :: ByteString
sCHEMA_KEY = "schema"

-- | Stores the units that are excluded (or included) from the base DB
uNITS_KEY :: ByteString
uNITS_KEY = "units"

-- | Stores the slices for the base DBs in a stack
sLICES_KEY :: ByteString
sLICES_KEY = "slices"

storeSchema :: Storage s => Database s -> StoredSchema -> IO ()
storeSchema db = Storage.store db sCHEMA_KEY . serializeCompact

retrieveSchema :: Storage s => Repo -> Database s -> IO (Maybe StoredSchema)
retrieveSchema repo db = do
  value <- Storage.retrieve db sCHEMA_KEY
  case deserializeCompact <$> value of
    Just (Right info) -> return $ Just info
    Just (Left msg) -> dbError repo $ "invalid schema: " ++ msg
    Nothing -> return Nothing

storeUnits :: Storage s => Database s -> [ByteString] -> IO ()
storeUnits db = Storage.store db uNITS_KEY . toStrict . encode

retrieveUnits :: Storage s => Repo -> Database s -> IO (Maybe [ByteString])
retrieveUnits repo db = do
  value <- Storage.retrieve db uNITS_KEY
  case decodeOrFail . fromStrict <$> value of
    Just (Right (_, _, units)) -> return $ Just units
    Just (Left (_, _, msg)) -> dbError repo $ "invalid units: " ++ msg
    Nothing -> return Nothing

-- Slices are each serialized using the RTS Slice::serialize(), and then
-- the list of serialized slices :: [ByteString] is serialized with
-- the Haskell Binary encoder.
storeSlices :: Storage s => Database s -> [Slice] -> IO ()
storeSlices db slices = do
  bytestrings <- mapM serializeSlice slices
  Storage.store db sLICES_KEY $ toStrict $ encode bytestrings

retrieveSlices :: Storage s => Repo -> Database s -> IO (Maybe [Slice])
retrieveSlices repo db = do
  value <- Storage.retrieve db sLICES_KEY
  case decodeOrFail . fromStrict <$> value of
    Just (Right (_, _, bytestrings)) ->
      Just <$> mapM deserializeSlice bytestrings
    Just (Left (_, _, msg)) -> dbError repo $ "invalid slices: " ++ msg
    Nothing -> return Nothing