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

-- | Compact public fact authority and exact contextual authorization.
module Felix.Checking.Authority
    ( EscapeKind(..)
    , EscapeKinds
    , emptyEscapeKinds
    , singletonEscapeKind
    , escapeKinds
    , escapeKindsToList
    , unionEscapeKinds
    , AuthoritySafety
    , cleanAuthoritySafety
    , authoritySafety
    , authoritySafetyEscapeKinds
    , unionAuthoritySafety
    , FactAuthority
    , factAuthority
    , factAuthorityTheorem
    , factAuthoritySafety
    , PreparedRequestDialect(..)
    , PreparedRequestMode(..)
    , PreparedRequestId
    , preparedRequestId
    , GuardedRuleSet
    , guardedRuleSet
    , guardedRuleTags
    , KernelConstructionDescriptor(..)
    , DatatypeCompilationDescriptor
    , datatypeCompilationDescriptor
    , CompilationDescriptor(..)
    , DirectAuthorization(..)
    , ValidationCertificate
    , validationCertificate
    , validationTarget
    , validationDirectAuthorization
    , ValidationCertificateError(..)
    , CandidateSafety
    , initialCandidateSafety
    , candidateSafetyAuthority
    , candidateFactAuthority
    , accumulateFactSafety
    , addCandidateEscape
    , FactSafetyError(..)
    , putEscapeKindsCache
    , getEscapeKindsCache
    , putAuthoritySafetyCache
    , getAuthoritySafetyCache
    , putFactAuthorityCache
    , getFactAuthorityCache
    , putPreparedRequestIdCache
    , getPreparedRequestIdCache
    , putDirectAuthorizationCache
    , getDirectAuthorizationCache
    , putValidationCertificateCache
    , getValidationCertificateCache
    ) where

import Base
import Felix.Checking.Foundation
import Felix.Checking.Identity
import Felix.Cache.Codec

import Control.DeepSeq (NFData)
import Control.Monad (unless, when)
import Data.Bits ((.&.), (.|.), bit, complement)
import Data.ByteString (ByteString)
import Data.ByteString qualified as ByteString
import Data.List (filter)
import Data.List.NonEmpty qualified as NonEmpty
import Data.Set qualified as Set
import Data.Word (Word8)


data EscapeKind
    = SourceAxiom
    | Omitted
    deriving stock (Show, Eq, Ord, Enum, Bounded, Generic)
    deriving anyclass (NFData)

-- | Canonical bounded set. The bit positions are stable cache tags.
newtype EscapeKinds =
    EscapeKinds Word8
    deriving stock (Show, Eq, Ord, Generic)
    deriving newtype (NFData)

emptyEscapeKinds :: EscapeKinds
emptyEscapeKinds =
    EscapeKinds 0

singletonEscapeKind :: EscapeKind -> EscapeKinds
singletonEscapeKind =
    EscapeKinds . escapeKindBit

escapeKinds :: Foldable collection => collection EscapeKind -> EscapeKinds
escapeKinds =
    foldl'
        (\acc kind ->
            unionEscapeKinds acc (singletonEscapeKind kind))
        emptyEscapeKinds

escapeKindsToList :: EscapeKinds -> [EscapeKind]
escapeKindsToList kinds =
    filter
        (\kind ->
            let EscapeKinds bits = kinds
            in bits .&. escapeKindBit kind /= 0)
        [SourceAxiom, Omitted]

unionEscapeKinds :: EscapeKinds -> EscapeKinds -> EscapeKinds
unionEscapeKinds (EscapeKinds left) (EscapeKinds right) =
    EscapeKinds (left .|. right)

escapeKindBit :: EscapeKind -> Word8
escapeKindBit = \case
    SourceAxiom -> bit 0
    Omitted -> bit 1

allEscapeBits :: Word8
allEscapeBits =
    escapeKindBit SourceAxiom .|. escapeKindBit Omitted


data AuthoritySafety
    = Clean
    | EscapeHatchBacked !EscapeKinds
    deriving stock (Show, Eq, Ord, Generic)
    deriving anyclass (NFData)

cleanAuthoritySafety :: AuthoritySafety
cleanAuthoritySafety =
    Clean

authoritySafety :: EscapeKinds -> AuthoritySafety
authoritySafety kinds
    | kinds == emptyEscapeKinds = Clean
    | otherwise = EscapeHatchBacked kinds

authoritySafetyEscapeKinds :: AuthoritySafety -> EscapeKinds
authoritySafetyEscapeKinds = \case
    Clean -> emptyEscapeKinds
    EscapeHatchBacked kinds -> kinds

unionAuthoritySafety
    :: AuthoritySafety
    -> AuthoritySafety
    -> AuthoritySafety
unionAuthoritySafety left right =
    authoritySafety
        (unionEscapeKinds
            (authoritySafetyEscapeKinds left)
            (authoritySafetyEscapeKinds right))


data FactAuthority = FactAuthority
    !TheoremRef
    !AuthoritySafety
    deriving stock (Show, Eq, Ord, Generic)
    deriving anyclass (NFData)

factAuthority :: TheoremRef -> AuthoritySafety -> FactAuthority
factAuthority =
    FactAuthority

factAuthorityTheorem :: FactAuthority -> TheoremRef
factAuthorityTheorem (FactAuthority reference _) =
    reference

factAuthoritySafety :: FactAuthority -> AuthoritySafety
factAuthoritySafety (FactAuthority _ safety) =
    safety


data PreparedRequestDialect
    = PreparedRequestFof
    | PreparedRequestTh0
    deriving stock (Show, Eq, Ord, Generic)

data PreparedRequestMode
    = PreparedRequestDirect
    | PreparedRequestIndirect
    deriving stock (Show, Eq, Ord, Generic)

newtype PreparedRequestId =
    PreparedRequestId CacheDigest
    deriving stock (Show, Eq, Ord, Generic)
    deriving newtype (Hashable, NFData)

preparedRequestId
    :: PreparedRequestDialect
    -> PreparedRequestMode
    -> ByteString
    -> PreparedRequestId
preparedRequestId dialect mode bytes =
    PreparedRequestId
        (hashCacheFields
            "felix-prepared-request-v2"
            [ ByteString.singleton (preparedRequestDialectTag dialect)
            , ByteString.singleton (preparedRequestModeTag mode)
            , bytes
            ])

preparedRequestDialectTag :: PreparedRequestDialect -> Word8
preparedRequestDialectTag = \case
    PreparedRequestFof -> 0x00
    PreparedRequestTh0 -> 0x01

preparedRequestModeTag :: PreparedRequestMode -> Word8
preparedRequestModeTag = \case
    PreparedRequestDirect -> 0x00
    PreparedRequestIndirect -> 0x01


-- | A nonempty canonical set of guarded kernel rules.
newtype GuardedRuleSet =
    GuardedRuleSet (Set KernelRuleTag)
    deriving stock (Show, Eq, Ord, Generic)

guardedRuleSet :: NonEmpty KernelRuleTag -> GuardedRuleSet
guardedRuleSet =
    GuardedRuleSet . Set.fromList . NonEmpty.toList

guardedRuleTags :: GuardedRuleSet -> Set KernelRuleTag
guardedRuleTags (GuardedRuleSet rules) =
    rules

data KernelConstructionDescriptor
    = FoundationLeaf !FoundationAxiomTag
    | GuardedFoundationRules !GuardedRuleSet
    | CheckedDefinitionEquation !ObjectId
    | CheckedSetConstructionExtensionality !ObjectId !CacheDigest
    deriving stock (Show, Eq, Ord, Generic)

-- | Exact semantic members of one trusted datatype compilation.
data DatatypeCompilationDescriptor =
    DatatypeCompilationDescriptor
        !ObjectId
        !(NonEmpty ObjectId)
        ![TheoremRef]
    deriving stock (Show, Eq, Ord, Generic)

datatypeCompilationDescriptor
    :: ObjectId
    -> NonEmpty ObjectId
    -> [TheoremRef]
    -> DatatypeCompilationDescriptor
datatypeCompilationDescriptor =
    DatatypeCompilationDescriptor

data CompilationDescriptor
    = DatatypeCompilation !DatatypeCompilationDescriptor
    deriving stock (Show, Eq, Ord, Generic)

data DirectAuthorization
    = CheckedKernelConstruction !KernelConstructionDescriptor
    | CheckedSourceProof ![PreparedRequestId]
    | TrustedCompilation !CompilationDescriptor
    | SourceAxiomAuthorization
    | OmittedAuthorization
    deriving stock (Show, Eq, Ord, Generic)


data ValidationCertificate = ValidationCertificate
    !FactAuthority
    !DirectAuthorization
    deriving stock (Show, Eq, Ord, Generic)

data ValidationCertificateError
    = SourceAxiomSafetyMismatch !AuthoritySafety
    | OmittedSafetyMissing !AuthoritySafety
    deriving stock (Show, Eq)

validationCertificate
    :: FactAuthority
    -> DirectAuthorization
    -> Either ValidationCertificateError ValidationCertificate
validationCertificate target direct = do
    case direct of
        SourceAxiomAuthorization ->
            unless
                (factAuthoritySafety target
                    == authoritySafety
                        (singletonEscapeKind SourceAxiom))
                (Left
                    (SourceAxiomSafetyMismatch
                        (factAuthoritySafety target)))
        OmittedAuthorization ->
            unless
                (Omitted
                    `elem` escapeKindsToList
                        (authoritySafetyEscapeKinds
                            (factAuthoritySafety target)))
                (Left
                    (OmittedSafetyMissing
                        (factAuthoritySafety target)))
        CheckedKernelConstruction{} -> pure ()
        CheckedSourceProof{} -> pure ()
        TrustedCompilation{} -> pure ()
    pure (ValidationCertificate target direct)

validationTarget :: ValidationCertificate -> FactAuthority
validationTarget (ValidationCertificate target _) =
    target

validationDirectAuthorization
    :: ValidationCertificate
    -> DirectAuthorization
validationDirectAuthorization (ValidationCertificate _ direct) =
    direct


newtype CandidateSafety =
    CandidateSafety AuthoritySafety
    deriving stock (Show, Eq, Ord, Generic)
    deriving newtype (NFData)

initialCandidateSafety :: CandidateSafety
initialCandidateSafety =
    CandidateSafety Clean

candidateSafetyAuthority :: CandidateSafety -> AuthoritySafety
candidateSafetyAuthority (CandidateSafety safety) =
    safety

-- | Freeze the final safety of one completed candidate into public authority.
--
-- The trusted completion boundary uses this projection for both the inert
-- validation certificate and its runtime pending authorization.
candidateFactAuthority
    :: TheoremRef
    -> CandidateSafety
    -> FactAuthority
candidateFactAuthority reference (CandidateSafety safety) =
    FactAuthority reference safety

data FactSafetyError
    = FactSafetyTheoremMismatch !TheoremRef !TheoremRef
    deriving stock (Show, Eq)

-- | Check theorem agreement and accumulate already-authorized public safety.
--
-- This operation does not authorize a freely constructed 'FactAuthority'.
-- Callers may use it only after checking opaque global, imported, or staged
-- builder authorization at the boundary that owns that authority.
accumulateFactSafety
    :: TheoremRef
    -> FactAuthority
    -> CandidateSafety
    -> Either FactSafetyError CandidateSafety
accumulateFactSafety expected supplied (CandidateSafety current) = do
    unless
        (factAuthorityTheorem supplied == expected)
        (Left
            (FactSafetyTheoremMismatch
                expected
                (factAuthorityTheorem supplied)))
    pure
        (CandidateSafety
            (unionAuthoritySafety
                current
                (factAuthoritySafety supplied)))

addCandidateEscape
    :: EscapeKind
    -> CandidateSafety
    -> CandidateSafety
addCandidateEscape kind (CandidateSafety current) =
    CandidateSafety
        (unionAuthoritySafety
            current
            (authoritySafety (singletonEscapeKind kind)))


putEscapeKindsCache :: EscapeKinds -> CachePut
putEscapeKindsCache (EscapeKinds bits) =
    putCacheTag bits

getEscapeKindsCache :: CacheGet EscapeKinds
getEscapeKindsCache = do
    bits <- getCacheTag
    unless
        (bits .&. complement allEscapeBits == 0)
        (fail "unknown cache escape-kind bit")
    pure (EscapeKinds bits)

putAuthoritySafetyCache :: AuthoritySafety -> CachePut
putAuthoritySafetyCache = \case
    Clean ->
        putCacheTag 0x00
    EscapeHatchBacked kinds -> do
        putCacheTag 0x01
        putEscapeKindsCache kinds

getAuthoritySafetyCache :: CacheGet AuthoritySafety
getAuthoritySafetyCache =
    getCacheTag >>= \case
        0x00 ->
            pure Clean
        0x01 -> do
            kinds <- getEscapeKindsCache
            when
                (kinds == emptyEscapeKinds)
                (fail "empty escape-hatch-backed safety")
            pure (EscapeHatchBacked kinds)
        tag ->
            fail ("unknown cache authority-safety tag " <> show tag)

putFactAuthorityCache :: FactAuthority -> CachePut
putFactAuthorityCache (FactAuthority reference safety) = do
    putTheoremRefCache reference
    putAuthoritySafetyCache safety

getFactAuthorityCache :: CacheGet FactAuthority
getFactAuthorityCache =
    FactAuthority
        <$> getTheoremRefCache
        <*> getAuthoritySafetyCache

putPreparedRequestIdCache :: PreparedRequestId -> CachePut
putPreparedRequestIdCache (PreparedRequestId digest) =
    putCacheDigest digest

getPreparedRequestIdCache :: CacheGet PreparedRequestId
getPreparedRequestIdCache =
    PreparedRequestId <$> getCacheDigest

putDirectAuthorizationCache :: DirectAuthorization -> CachePut
putDirectAuthorizationCache = \case
    CheckedKernelConstruction descriptor -> do
        putCacheTag 0x00
        putKernelDescriptor descriptor
    CheckedSourceProof requests -> do
        putCacheTag 0x01
        putCacheList putPreparedRequestIdCache requests
    TrustedCompilation descriptor -> do
        putCacheTag 0x02
        putCompilationDescriptor descriptor
    SourceAxiomAuthorization ->
        putCacheTag 0x03
    OmittedAuthorization ->
        putCacheTag 0x04

getDirectAuthorizationCache :: CacheGet DirectAuthorization
getDirectAuthorizationCache =
    getCacheTag >>= \case
        0x00 ->
            CheckedKernelConstruction <$> getKernelDescriptor
        0x01 ->
            CheckedSourceProof
                <$> getCacheList getPreparedRequestIdCache
        0x02 ->
            TrustedCompilation <$> getCompilationDescriptor
        0x03 ->
            pure SourceAxiomAuthorization
        0x04 ->
            pure OmittedAuthorization
        tag ->
            fail ("unknown cache direct-authorization tag " <> show tag)

putValidationCertificateCache :: ValidationCertificate -> CachePut
putValidationCertificateCache
        (ValidationCertificate target direct) = do
    putFactAuthorityCache target
    putDirectAuthorizationCache direct

getValidationCertificateCache :: CacheGet ValidationCertificate
getValidationCertificateCache = do
    target <- getFactAuthorityCache
    direct <- getDirectAuthorizationCache
    case validationCertificate target direct of
        Left err ->
            fail ("invalid cache validation certificate: " <> show err)
        Right certificate ->
            pure certificate


putKernelDescriptor :: KernelConstructionDescriptor -> CachePut
putKernelDescriptor = \case
    FoundationLeaf tag -> do
        putCacheTag 0x00
        putFoundationTag tag
    GuardedFoundationRules rules -> do
        putCacheTag 0x01
        putCacheList putRuleTag
            (Set.toAscList (guardedRuleTags rules))
    CheckedDefinitionEquation identity -> do
        putCacheTag 0x02
        putObjectIdCache identity
    CheckedSetConstructionExtensionality identity construction -> do
        putCacheTag 0x03
        putObjectIdCache identity
        putCacheDigest construction

getKernelDescriptor :: CacheGet KernelConstructionDescriptor
getKernelDescriptor =
    getCacheTag >>= \case
        0x00 -> FoundationLeaf <$> getFoundationTag
        0x01 -> do
            tags <- getCacheList getRuleTag
            rules <-
                maybe
                    (fail "cache guarded-rule set is empty")
                    pure
                    (NonEmpty.nonEmpty tags)
            unless (strictlyIncreasing tags)
                (fail "cache guarded-rule tags are not in canonical order")
            pure (GuardedFoundationRules (guardedRuleSet rules))
        0x02 -> CheckedDefinitionEquation <$> getObjectIdCache
        0x03 ->
            CheckedSetConstructionExtensionality
                <$> getObjectIdCache
                <*> getCacheDigest
        tag ->
            fail ("unknown cache kernel-construction tag " <> show tag)

putCompilationDescriptor :: CompilationDescriptor -> CachePut
putCompilationDescriptor
        (DatatypeCompilation
            (DatatypeCompilationDescriptor
                carrier constructors facts)) = do
    putCacheTag 0x00
    putObjectIdCache carrier
    putCacheList putObjectIdCache (NonEmpty.toList constructors)
    putCacheList putTheoremRefCache facts

getCompilationDescriptor :: CacheGet CompilationDescriptor
getCompilationDescriptor =
    getCacheTag >>= \case
        0x00 -> do
            carrier <- getObjectIdCache
            rawConstructors <-
                getCacheList getObjectIdCache
            constructors <-
                maybe
                    (fail "cache datatype compilation has no constructors")
                    pure
                    (NonEmpty.nonEmpty rawConstructors)
            facts <- getCacheList getTheoremRefCache
            pure
                (DatatypeCompilation
                    (DatatypeCompilationDescriptor
                        carrier constructors facts))
        tag ->
            fail ("unknown cache compilation tag " <> show tag)

putFoundationTag :: FoundationAxiomTag -> CachePut
putFoundationTag =
    putCacheTag . \case
        EmptyCharacteristic -> 0x00
        PairSetCharacteristic -> 0x01
        FamilyUnionCharacteristic -> 0x02
        PowerSetCharacteristic -> 0x03
        SeparationCharacteristic -> 0x04
        ReplacementCharacteristic -> 0x05
        SetChooseWitness -> 0x06
        SetExtensionality -> 0x07
        SetInduction -> 0x08
        PropositionalExtensionality -> 0x09
        DoubleNegationElim -> 0x0a
        UnivOfContains -> 0x0b
        UnivOfTransitive -> 0x0c
        UnivOfFamilyUnionClosed -> 0x0d
        UnivOfPowerSetClosed -> 0x0e
        UnivOfReplacementClosed -> 0x0f
        UnivOfMinimal -> 0x10

getFoundationTag :: CacheGet FoundationAxiomTag
getFoundationTag =
    getCacheTag >>= \case
        0x00 -> pure EmptyCharacteristic
        0x01 -> pure PairSetCharacteristic
        0x02 -> pure FamilyUnionCharacteristic
        0x03 -> pure PowerSetCharacteristic
        0x04 -> pure SeparationCharacteristic
        0x05 -> pure ReplacementCharacteristic
        0x06 -> pure SetChooseWitness
        0x07 -> pure SetExtensionality
        0x08 -> pure SetInduction
        0x09 -> pure PropositionalExtensionality
        0x0a -> pure DoubleNegationElim
        0x0b -> pure UnivOfContains
        0x0c -> pure UnivOfTransitive
        0x0d -> pure UnivOfFamilyUnionClosed
        0x0e -> pure UnivOfPowerSetClosed
        0x0f -> pure UnivOfReplacementClosed
        0x10 -> pure UnivOfMinimal
        tag ->
            fail ("unknown cache foundation-axiom tag " <> show tag)

putRuleTag :: KernelRuleTag -> CachePut
putRuleTag =
    putCacheTag . \case
        SetLfpBound -> 0x00
        SetLfpLeast -> 0x01
        SetLfpFixed -> 0x02
        SetLfpInduct -> 0x03

getRuleTag :: CacheGet KernelRuleTag
getRuleTag =
    getCacheTag >>= \case
        0x00 -> pure SetLfpBound
        0x01 -> pure SetLfpLeast
        0x02 -> pure SetLfpFixed
        0x03 -> pure SetLfpInduct
        tag ->
            fail ("unknown cache kernel-rule tag " <> show tag)

strictlyIncreasing :: Ord value => [value] -> Bool
strictlyIncreasing values =
    and (zipWith (<) values (drop 1 values))