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

-- | Stable byte primitives used only by durable mathematical identities.
module Felix.Math.Codec
    ( MathematicalDigest
    , mathematicalDigestBytes
    , mathematicalDigestFromBytes
    , mathematicalDigestHex
    , hashCanonicalFields
    , encodeU32
    , encodeU64
    , encodeFrame
    , encodeNatural
    , encodeInteger
    , encodeSequence
    , encodeCanonicalPathRecord
    , encodeCanonicalSet
    , encodeCanonicalMap
    , encodeCoreType
    , encodeCoreIntrinsicTag
    , encodeCanonicalTerm
    , MathematicalCodecError(..)
    ) where

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

import Control.DeepSeq (NFData)
import Crypto.Hash qualified as Crypto
import Data.Bits ((.&.), shiftR)
import Data.ByteArray qualified as ByteArray
import Data.ByteString (ByteString)
import Data.ByteString qualified as ByteString
import Data.ByteString.Builder qualified as Builder
import Data.ByteString.Lazy qualified as LazyByteString
import Data.List qualified as List
import Data.Text qualified as Text
import Data.Text.Encoding qualified as Text
import Data.Word (Word32)
import Numeric.Natural (Natural)


-- | One raw SHA-256 digest. Hexadecimal is presentation only.
newtype MathematicalDigest =
    MathematicalDigest ByteString
    deriving stock (Eq, Ord, Generic)
    deriving newtype (Hashable, NFData)

instance Show MathematicalDigest where
    show =
        Text.unpack . mathematicalDigestHex

mathematicalDigestBytes :: MathematicalDigest -> ByteString
mathematicalDigestBytes (MathematicalDigest bytes) =
    bytes

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

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

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


data MathematicalCodecError
    = CanonicalFieldCountOverflow !Integer
    | CanonicalSequenceCountOverflow !Integer
    | CanonicalPathComponentCountOverflow !Integer
    | DuplicateCanonicalSetElement !ByteString
    | DuplicateCanonicalMapKey !ByteString
    deriving stock (Show, Eq)

-- | The sole durable hash framing primitive.
hashCanonicalFields
    :: Text
    -> [ByteString]
    -> Either MathematicalCodecError MathematicalDigest
hashCanonicalFields domain fields = do
    fieldCount <- checkedU32
        CanonicalFieldCountOverflow
        (length fields)
    let input =
            encodeFrame (Text.encodeUtf8 domain)
                <> encodeU32 fieldCount
                <> foldMap encodeFrame fields
        digest =
            Crypto.hash input
                :: Crypto.Digest Crypto.SHA256
    pure
        (MathematicalDigest
            (ByteArray.convert digest))

encodeU32 :: Word32 -> ByteString
encodeU32 =
    strictBuilder . Builder.word32BE

encodeU64 :: Word64 -> ByteString
encodeU64 =
    strictBuilder . Builder.word64BE

encodeFrame :: ByteString -> ByteString
encodeFrame bytes =
    encodeU64 (fromIntegral (ByteString.length bytes))
        <> bytes

-- | Minimal unsigned big-endian magnitude, framed by its byte length.
encodeNatural :: Natural -> ByteString
encodeNatural number =
    encodeU64 (fromIntegral (ByteString.length magnitude))
        <> magnitude
  where
    magnitude
        | number == 0 =
            ByteString.empty
        | otherwise =
            ByteString.reverse
                (ByteString.unfoldr step number)

    step 0 =
        Nothing
    step remaining =
        Just
            ( fromIntegral (remaining .&. 0xff)
            , remaining `shiftR` 8
            )

encodeInteger :: Integer -> ByteString
encodeInteger integer
    | integer < 0 =
        ByteString.cons 0x01
            (encodeNatural (fromInteger (negate integer)))
    | otherwise =
        ByteString.cons 0x00
            (encodeNatural (fromInteger integer))

encodeSequence
    :: [ByteString]
    -> Either MathematicalCodecError ByteString
encodeSequence elements = do
    elementCount <- checkedU64
        CanonicalSequenceCountOverflow
        (length elements)
    pure
        (encodeU64 elementCount
            <> foldMap encodeFrame elements)

-- | Encode an already validated path-component vector under its exact domain.
encodeCanonicalPathRecord
    :: Text
    -> [Text]
    -> Either MathematicalCodecError ByteString
encodeCanonicalPathRecord domain components = do
    componentCount <- checkedU32
        CanonicalPathComponentCountOverflow
        (length components)
    pure
        (encodeFrame (Text.encodeUtf8 domain)
            <> encodeU32 componentCount
            <> foldMap
                (encodeFrame . Text.encodeUtf8)
                components)

encodeCanonicalSet
    :: [ByteString]
    -> Either MathematicalCodecError ByteString
encodeCanonicalSet elements = do
    let ordered = List.sort elements
    case firstAdjacentDuplicate ordered of
        Just duplicate ->
            Left (DuplicateCanonicalSetElement duplicate)
        Nothing ->
            encodeSequence ordered

encodeCanonicalMap
    :: [(ByteString, ByteString)]
    -> Either MathematicalCodecError ByteString
encodeCanonicalMap entries = do
    let ordered =
            List.sortOn fst entries
    case firstAdjacentDuplicate (fst <$> ordered) of
        Just duplicate ->
            Left (DuplicateCanonicalMapKey duplicate)
        Nothing ->
            encodeSequence
                [ encodeFrame key <> encodeFrame value
                | (key, value) <- ordered
                ]

encodeCoreType :: CoreType -> ByteString
encodeCoreType = \case
    TyProp ->
        ByteString.singleton 0x00
    TySet ->
        ByteString.singleton 0x01
    TyArrow domain codomain ->
        ByteString.singleton 0x02
            <> encodeFrame (encodeCoreType domain)
            <> encodeFrame (encodeCoreType codomain)

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

encodeCanonicalTerm
    :: (global -> ByteString)
    -> CanonicalTerm global
    -> ByteString
encodeCanonicalTerm encodeGlobal = \case
    CBound index ->
        ByteString.singleton 0x00
            <> encodeNatural index
    CGlobal global ->
        ByteString.singleton 0x01
            <> encodeFrame (encodeGlobal global)
    CIntrinsic intrinsic ->
        ByteString.singleton 0x02
            <> encodeCoreIntrinsicTag intrinsic
    COpaqueInteger integer ->
        ByteString.singleton 0x03
            <> encodeInteger integer
    CApp function argument ->
        ByteString.singleton 0x04
            <> encodeFrame
                (encodeCanonicalTerm encodeGlobal function)
            <> encodeFrame
                (encodeCanonicalTerm encodeGlobal argument)
    CLam binderType body ->
        ByteString.singleton 0x05
            <> encodeFrame (encodeCoreType binderType)
            <> encodeFrame
                (encodeCanonicalTerm encodeGlobal body)
    CFalsum ->
        ByteString.singleton 0x06
    CImp premise conclusion ->
        ByteString.singleton 0x07
            <> encodeFrame
                (encodeCanonicalTerm encodeGlobal premise)
            <> encodeFrame
                (encodeCanonicalTerm encodeGlobal conclusion)
    CEq operandType left right ->
        ByteString.singleton 0x08
            <> encodeFrame (encodeCoreType operandType)
            <> encodeFrame
                (encodeCanonicalTerm encodeGlobal left)
            <> encodeFrame
                (encodeCanonicalTerm encodeGlobal right)
    CForall binderType body ->
        ByteString.singleton 0x09
            <> encodeFrame (encodeCoreType binderType)
            <> encodeFrame
                (encodeCanonicalTerm encodeGlobal body)

strictBuilder :: Builder.Builder -> ByteString
strictBuilder =
    LazyByteString.toStrict . Builder.toLazyByteString

checkedU32
    :: (Integer -> MathematicalCodecError)
    -> Int
    -> Either MathematicalCodecError Word32
checkedU32 makeError suppliedCount
    | suppliedCountInteger > toInteger (maxBound :: Word32) =
        Left (makeError suppliedCountInteger)
    | otherwise =
        Right (fromIntegral suppliedCount)
  where
    suppliedCountInteger = toInteger suppliedCount

checkedU64
    :: (Integer -> MathematicalCodecError)
    -> Int
    -> Either MathematicalCodecError Word64
checkedU64 makeError suppliedCount
    | suppliedCountInteger > toInteger (maxBound :: Word64) =
        Left (makeError suppliedCountInteger)
    | otherwise =
        Right (fromIntegral suppliedCount)
  where
    suppliedCountInteger = toInteger suppliedCount

firstAdjacentDuplicate :: Eq a => [a] -> Maybe a
firstAdjacentDuplicate = \case
    left : right : rest
        | left == right ->
            Just left
        | otherwise ->
            firstAdjacentDuplicate (right : rest)
    _ ->
        Nothing