{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE NoImplicitPrelude #-}

-- | Deterministic, checked serialization scoped only by 'CacheEpoch'.
--
-- These bytes are disposable implementation data. They are never input to a
-- durable mathematical identity.
module Felix.Cache.Codec
    ( CacheEpoch
    , currentCacheEpoch
    , cacheEpochValue
    , cacheEpochFromValue
    , CacheDigest
    , cacheDigestBytes
    , cacheDigestFromBytes
    , cacheDigestHex
    , hashCacheFields
    , CachePut
    , CacheGet
    , CacheDecodeError(..)
    , encodeCache
    , decodeCache
    , putCacheTag
    , getCacheTag
    , putCacheBytes
    , getCacheBytes
    , putCacheText
    , getCacheText
    , putCacheNatural
    , getCacheNatural
    , putCacheInteger
    , getCacheInteger
    , putMathematicalDigestCache
    , getMathematicalDigestCache
    , putCacheDigest
    , getCacheDigest
    , putCacheList
    , getCacheList
    , putCacheMaybe
    , getCacheMaybe
    , putCanonicalCacheMap
    , getCanonicalCacheMap
    , putCoreTypeCache
    , getCoreTypeCache
    , putCoreIntrinsicTagCache
    , getCoreIntrinsicTagCache
    , putCanonicalTermCache
    , getCanonicalTermCache
    ) where

import Base hiding (Empty)
import Felix.Checking.Core
import Felix.Math.Codec

import Control.DeepSeq (NFData)
import Control.Monad (replicateM, unless)
import Crypto.Hash qualified as Crypto
import Data.Binary.Get qualified as Get
import Data.Binary.Put qualified as Put
import Data.Bits ((.&.), (.|.), shiftL, shiftR)
import Data.ByteArray qualified as ByteArray
import Data.ByteString (ByteString)
import Data.ByteString qualified as ByteString
import Data.ByteString.Lazy qualified as LazyByteString
import Data.List qualified as List
import Data.Map.Strict qualified as Map
import Data.Text qualified as StrictText
import Data.Text.Encoding qualified as Text
import Data.Word (Word8, Word32)
import Numeric.Natural (Natural)


newtype CacheEpoch = CacheEpoch Word32
    deriving stock (Show, Eq, Ord)

currentCacheEpoch :: CacheEpoch
currentCacheEpoch =
    CacheEpoch 38

cacheEpochValue :: CacheEpoch -> Word32
cacheEpochValue (CacheEpoch value) =
    value

cacheEpochFromValue :: Word32 -> CacheEpoch
cacheEpochFromValue =
    CacheEpoch

-- | A SHA-256 digest whose meaning is scoped by the store-wide cache epoch.
-- Its input encoding is deliberately separate from mathematical identity.
newtype CacheDigest =
    CacheDigest ByteString
    deriving stock (Eq, Ord, Generic)
    deriving newtype (Hashable, NFData)

instance Show CacheDigest where
    show =
        StrictText.unpack . cacheDigestHex

cacheDigestBytes :: CacheDigest -> ByteString
cacheDigestBytes (CacheDigest bytes) =
    bytes

cacheDigestFromBytes :: ByteString -> Maybe CacheDigest
cacheDigestFromBytes bytes
    | ByteString.length bytes == 32 =
        Just (CacheDigest bytes)
    | otherwise =
        Nothing

cacheDigestHex :: CacheDigest -> Text
cacheDigestHex =
    Text.decodeUtf8
        . ByteString.concatMap encodeHexByte
        . cacheDigestBytes
  where
    encodeHexByte byte =
        ByteString.pack
            [ hexDigit (byte `shiftR` 4)
            , hexDigit (byte .&. 0x0f)
            ]

    hexDigit nibble
        | nibble < 10 =
            48 + nibble
        | otherwise =
            87 + nibble

-- | Hash deterministic cache fields under a cache-only domain. These bytes
-- may change only together with 'currentCacheEpoch'.
hashCacheFields :: Text -> [ByteString] -> CacheDigest
hashCacheFields domain fields =
    CacheDigest
        (ByteArray.convert
            (Crypto.hash input
                :: Crypto.Digest Crypto.SHA256))
  where
    input =
        encodeCache do
            putCacheText domain
            putCacheList putCacheBytes fields

type CachePut = Put.Put
type CacheGet = Get.Get

data CacheDecodeError
    = CacheBinaryDecodeError !String
    | CacheTrailingBytes !Word64
    deriving stock (Show, Eq)

encodeCache :: CachePut -> ByteString
encodeCache =
    LazyByteString.toStrict . Put.runPut

decodeCache
    :: CacheGet value
    -> ByteString
    -> Either CacheDecodeError value
decodeCache decoder bytes =
    case Get.runGetOrFail decoder (LazyByteString.fromStrict bytes) of
        Left (_remaining, _offset, message) ->
            Left (CacheBinaryDecodeError message)
        Right (remaining, _offset, value)
            | LazyByteString.null remaining ->
                Right value
            | otherwise ->
                Left
                    (CacheTrailingBytes
                        (fromIntegral
                            (LazyByteString.length remaining)))

putCacheTag :: Word8 -> CachePut
putCacheTag =
    Put.putWord8

getCacheTag :: CacheGet Word8
getCacheTag =
    Get.getWord8

putCacheBytes :: ByteString -> CachePut
putCacheBytes bytes = do
    Put.putWord64be
        (fromIntegral (ByteString.length bytes))
    Put.putByteString bytes

getCacheBytes :: CacheGet ByteString
getCacheBytes = do
    byteCount <- Get.getWord64be
    checkedByteCount <-
        checkedCacheCount "byte string" byteCount
    Get.getByteString checkedByteCount

putCacheText :: Text -> CachePut
putCacheText =
    putCacheBytes . Text.encodeUtf8

getCacheText :: CacheGet Text
getCacheText = do
    bytes <- getCacheBytes
    case Text.decodeUtf8' bytes of
        Left err ->
            fail ("invalid cache UTF-8: " <> show err)
        Right text ->
            pure text

putCacheNatural :: Natural -> CachePut
putCacheNatural number =
    putCacheBytes (naturalMagnitude number)

getCacheNatural :: CacheGet Natural
getCacheNatural = do
    magnitude <- getCacheBytes
    case ByteString.uncons magnitude of
        Just (0, _) ->
            fail "non-minimal cache natural"
        _ ->
            pure
                (ByteString.foldl'
                    (\value byte ->
                        (value `shiftL` 8)
                            .|. fromIntegral byte)
                    0
                    magnitude)

putCacheInteger :: Integer -> CachePut
putCacheInteger integer
    | integer < 0 = do
        putCacheTag 0x01
        putCacheNatural (fromInteger (negate integer))
    | otherwise = do
        putCacheTag 0x00
        putCacheNatural (fromInteger integer)

getCacheInteger :: CacheGet Integer
getCacheInteger = do
    sign <- getCacheTag
    magnitude <- getCacheNatural
    case sign of
        0x00 ->
            pure (toInteger magnitude)
        0x01
            | magnitude == 0 ->
                fail "negative zero in cache integer"
            | otherwise ->
                pure (negate (toInteger magnitude))
        _ ->
            fail ("unknown cache integer sign " <> show sign)

putMathematicalDigestCache
    :: MathematicalDigest
    -> CachePut
putMathematicalDigestCache =
    Put.putByteString . mathematicalDigestBytes

getMathematicalDigestCache :: CacheGet MathematicalDigest
getMathematicalDigestCache = do
    bytes <- Get.getByteString 32
    case mathematicalDigestFromBytes bytes of
        Just digest ->
            pure digest
        Nothing ->
            fail "cache mathematical digest is not 32 bytes"

putCacheDigest :: CacheDigest -> CachePut
putCacheDigest =
    Put.putByteString . cacheDigestBytes

getCacheDigest :: CacheGet CacheDigest
getCacheDigest = do
    bytes <- Get.getByteString 32
    case cacheDigestFromBytes bytes of
        Just digest ->
            pure digest
        Nothing ->
            fail "cache digest is not 32 bytes"

putCacheList
    :: (value -> CachePut)
    -> [value]
    -> CachePut
putCacheList putValue values = do
    Put.putWord64be (fromIntegral (length values))
    traverse_ putValue values

getCacheList
    :: CacheGet value
    -> CacheGet [value]
getCacheList getValue = do
    rawCount <- Get.getWord64be
    valueCount <- checkedCacheCount "list" rawCount
    replicateM valueCount getValue

putCacheMaybe
    :: (value -> CachePut)
    -> Maybe value
    -> CachePut
putCacheMaybe putValue = \case
    Nothing ->
        putCacheTag 0x00
    Just value -> do
        putCacheTag 0x01
        putValue value

getCacheMaybe
    :: CacheGet value
    -> CacheGet (Maybe value)
getCacheMaybe getValue =
    getCacheTag >>= \case
        0x00 ->
            pure Nothing
        0x01 ->
            Just <$> getValue
        tag ->
            fail ("unknown cache maybe tag " <> show tag)

-- | Maps are ordered by encoded semantic key, independently of runtime-map
-- traversal.
putCanonicalCacheMap
    :: (key -> CachePut)
    -> (value -> CachePut)
    -> Map key value
    -> CachePut
putCanonicalCacheMap putKey putValue entries =
    putCacheList putEntry ordered
  where
    ordered =
        List.sortOn fst
            [ (encodeCache (putKey key), value)
            | (key, value) <- Map.toList entries
            ]

    putEntry (keyBytes, value) = do
        putCacheBytes keyBytes
        putCacheBytes (encodeCache (putValue value))

getCanonicalCacheMap
    :: Ord key
    => CacheGet key
    -> CacheGet value
    -> CacheGet (Map key value)
getCanonicalCacheMap getKey getValue = do
    encodedEntries <-
        getCacheList
            ((,) <$> getCacheBytes <*> getCacheBytes)
    verifyStrictlyAscending (fst <$> encodedEntries)
    entries <- traverse decodeEntry encodedEntries
    let result =
            Map.fromList entries
    unless
        (Map.size result == length entries)
        (fail "duplicate decoded cache map key")
    pure result
  where
    decodeEntry (keyBytes, valueBytes) =
        (,) <$> nested "map key" getKey keyBytes
            <*> nested "map value" getValue valueBytes

    nested description decoder bytes =
        case decodeCache decoder bytes of
            Left err ->
                fail
                    ("invalid cache "
                        <> description
                        <> ": "
                        <> show err)
            Right value ->
                pure value

verifyStrictlyAscending :: [ByteString] -> CacheGet ()
verifyStrictlyAscending = \case
    left : right : rest
        | left < right ->
            verifyStrictlyAscending (right : rest)
        | otherwise ->
            fail "cache map keys are not in canonical order"
    _ ->
        pure ()


putCoreTypeCache :: CoreType -> CachePut
putCoreTypeCache = \case
    TyProp ->
        putCacheTag 0x00
    TySet ->
        putCacheTag 0x01
    TyArrow domain codomain -> do
        putCacheTag 0x02
        putCoreTypeCache domain
        putCoreTypeCache codomain

getCoreTypeCache :: CacheGet CoreType
getCoreTypeCache =
    getCacheTag >>= \case
        0x00 ->
            pure TyProp
        0x01 ->
            pure TySet
        0x02 ->
            TyArrow
                <$> getCoreTypeCache
                <*> getCoreTypeCache
        tag ->
            fail ("unknown cache core-type tag " <> show tag)

putCoreIntrinsicTagCache :: CoreIntrinsicTag -> CachePut
putCoreIntrinsicTagCache =
    putCacheTag . \case
        Member ->
            0x00
        Empty ->
            0x01
        PairSet ->
            0x02
        FamilyUnion ->
            0x03
        PowerSet ->
            0x04
        Sep ->
            0x05
        Repl ->
            0x06
        SetChoose ->
            0x07
        UnivOf ->
            0x08
        ISetLfp ->
            0x09

getCoreIntrinsicTagCache :: CacheGet CoreIntrinsicTag
getCoreIntrinsicTagCache =
    getCacheTag >>= \case
        0x00 ->
            pure Member
        0x01 ->
            pure Empty
        0x02 ->
            pure PairSet
        0x03 ->
            pure FamilyUnion
        0x04 ->
            pure PowerSet
        0x05 ->
            pure Sep
        0x06 ->
            pure Repl
        0x07 ->
            pure SetChoose
        0x08 ->
            pure UnivOf
        0x09 ->
            pure ISetLfp
        tag ->
            fail ("unknown cache intrinsic tag " <> show tag)

putCanonicalTermCache
    :: (global -> CachePut)
    -> CanonicalTerm global
    -> CachePut
putCanonicalTermCache putGlobal = \case
    CBound index -> do
        putCacheTag 0x00
        putCacheNatural index
    CGlobal global -> do
        putCacheTag 0x01
        putGlobal global
    CIntrinsic intrinsic -> do
        putCacheTag 0x02
        putCoreIntrinsicTagCache intrinsic
    COpaqueInteger integer -> do
        putCacheTag 0x03
        putCacheInteger integer
    CApp function argument -> do
        putCacheTag 0x04
        putCanonicalTermCache putGlobal function
        putCanonicalTermCache putGlobal argument
    CLam binderType body -> do
        putCacheTag 0x05
        putCoreTypeCache binderType
        putCanonicalTermCache putGlobal body
    CFalsum ->
        putCacheTag 0x06
    CImp premise conclusion -> do
        putCacheTag 0x07
        putCanonicalTermCache putGlobal premise
        putCanonicalTermCache putGlobal conclusion
    CEq operandType left right -> do
        putCacheTag 0x08
        putCoreTypeCache operandType
        putCanonicalTermCache putGlobal left
        putCanonicalTermCache putGlobal right
    CForall binderType body -> do
        putCacheTag 0x09
        putCoreTypeCache binderType
        putCanonicalTermCache putGlobal body

getCanonicalTermCache
    :: CacheGet global
    -> CacheGet (CanonicalTerm global)
getCanonicalTermCache getGlobal =
    getCacheTag >>= \case
        0x00 ->
            CBound <$> getCacheNatural
        0x01 ->
            CGlobal <$> getGlobal
        0x02 ->
            CIntrinsic <$> getCoreIntrinsicTagCache
        0x03 ->
            COpaqueInteger <$> getCacheInteger
        0x04 ->
            CApp
                <$> getCanonicalTermCache getGlobal
                <*> getCanonicalTermCache getGlobal
        0x05 ->
            CLam
                <$> getCoreTypeCache
                <*> getCanonicalTermCache getGlobal
        0x06 ->
            pure CFalsum
        0x07 ->
            CImp
                <$> getCanonicalTermCache getGlobal
                <*> getCanonicalTermCache getGlobal
        0x08 ->
            CEq
                <$> getCoreTypeCache
                <*> getCanonicalTermCache getGlobal
                <*> getCanonicalTermCache getGlobal
        0x09 ->
            CForall
                <$> getCoreTypeCache
                <*> getCanonicalTermCache getGlobal
        tag ->
            fail ("unknown cache canonical-term tag " <> show tag)


naturalMagnitude :: Natural -> ByteString
naturalMagnitude number
    | number == 0 =
        ByteString.empty
    | otherwise =
        ByteString.reverse
            (ByteString.unfoldr step number)
  where
    step 0 =
        Nothing
    step remaining =
        Just
            ( fromIntegral (remaining `mod` 256)
            , remaining `div` 256
            )

checkedCacheCount :: String -> Word64 -> CacheGet Int
checkedCacheCount description supplied
    | supplied > fromIntegral (maxBound :: Int) =
        fail
            ("cache "
                <> description
                <> " length exceeds Int")
    | otherwise =
        pure (fromIntegral supplied)