module arch.system.identity.identity;

import prelude open;

type OrdKey OrdKeyType :=
  mkOrdkey@{
    compare : OrdKeyType -> OrdKeyType -> Ordering;
  };

type HASH OrdKeyType Hashable :=
  mkHASH@{
    ordKey : OrdKey OrdKeyType;
    hash : Hashable -> OrdKeyType;
  };

-- Note: instance of this with Data.Map should be made
type OrdMap (OrdKeyType : Type) (MapCon : Type -> Type) :=
  mkOrdMap@{
    ordKey : OrdKey OrdKeyType;
    empty {A} : MapCon A;
    map {A B} : (A -> B) -> MapCon A -> MapCon B;
    insert {A} : Pair (MapCon A) (Pair OrdKeyType A) -> MapCon A;
    foldl {A B} : (Pair A B -> B) -> B -> MapCon A -> B;
    intersectWith
      {A B C} : (Pair A B -> C) -> Pair (MapCon A) (MapCon B) -> MapCon C;
    all {A} : (A -> Bool) -> MapCon A -> Bool;
  -- Bunch of stuff, see https://www.smlnj.org/doc/smlnj-lib/Util/sig-OrdMap.html
  };

type Signer SignerType Signable Commitment :=
  mkSigner@{
    sign : SignerType -> Signable -> Commitment;
  };

type Decryptor DecryptorType Plaintext Ciphertext :=
  mkDecryptor@{
    decrypt : DecryptorType -> Ciphertext -> Option Plaintext;
  };

type InternalIdentity SignerType Signable Commitment DecryptorType Plaintext Ciphertext :=
  mkInternalIdentity@{
    signer : Signer SignerType Signable Commitment;
    decryptor : Decryptor DecryptorType Plaintext Ciphertext;
  };

type Verifier OrdKey VerifierType Signable Commitment :=
  mkVerifier@{
    verify : VerifierType -> Signable -> Commitment -> Bool;
    verifierHash : HASH OrdKey VerifierType;
  };

type Encryptor OrdKey EncryptorType Plaintext Ciphertext :=
  mkEncryptor@{
    encrypt : EncryptorType -> Plaintext -> Ciphertext;
    encryptorHash : HASH OrdKey EncryptorType;
  };

type ExternalIdentity VerifierOrdKeyType VerifierType Signable Commitment EncryptorOrdKeyType EncryptorType Plaintext Ciphertext :=
  mkExternalIdentity@{
    verifier : Verifier VerifierOrdKeyType VerifierType Signable Commitment;
    encryptor
      : Encryptor EncryptorOrdKeyType EncryptorType Plaintext Ciphertext;
  };

type Identity SignerType InternalSignable InternalCommitment DecryptorType InternalCiphertext InternalPlaintext VerifierOrdKeyType VerifierType ExternalSignable ExternalCommitment EncryptorOrdKeyType EncryptorType ExternalPlaintext ExternalCiphertext :=
  mkIdentity@{
    internalIdentity
      : InternalIdentity
        SignerType
        InternalSignable
        InternalCommitment
        DecryptorType
        InternalPlaintext
        InternalCiphertext;
    externalIdentity
      : ExternalIdentity
        VerifierOrdKeyType
        VerifierType
        ExternalSignable
        ExternalCommitment
        EncryptorOrdKeyType
        EncryptorType
        ExternalPlaintext
        ExternalCiphertext;
  };

type SignsFor OrdKey VerifierType Signable Commitment Evidence :=
  mkSignsFor@{
    verifier : Verifier OrdKey VerifierType Signable Commitment;
    signsFor : Evidence -> Pair VerifierType VerifierType -> Bool;
  };

type ReadsFor (OrdKey EncryptorType Plaintext Ciphertext Evidence : Type) :=
  mkReadsFor@{
    encryptor : Encryptor OrdKey EncryptorType Plaintext Ciphertext;
    readsFor : Evidence -> Pair EncryptorType EncryptorType -> Bool;
  };

type ComposeHashable (VerifierType : Type) (MapCon : Type -> Type) :=
  mkComposeHashable@{
    threshold : Nat;
    weights : MapCon (Pair Nat VerifierType);
  };

type ThresholdCompose (OrdKey : Type) (MapCon : Type
  -> Type) (VerifierType Signable Commitment SignerType VerifierHashOrdKeyType : Type) :=
  mkThresholdCompose@{
    map : OrdMap OrdKey MapCon;
    underlyingVerifier : Verifier OrdKey VerifierType Signable Commitment;
    underlyingSigner : Signer SignerType Signable Commitment;
    verifierHash
      : HASH VerifierHashOrdKeyType (ComposeHashable VerifierType MapCon);
    
    sign : MapCon SignerType -> Signable -> MapCon Commitment;
    verify
      : ComposeHashable VerifierType MapCon
        -> Signable
        -> MapCon Commitment
        -> Bool;
    signerCompose : List (Pair VerifierType SignerType) -> MapCon SignerType;
    verifierCompose
      : Nat
        -> List (Pair Nat VerifierType)
        -> ComposeHashable VerifierType MapCon;
    verifierAnd
      : VerifierType -> VerifierType -> ComposeHashable VerifierType MapCon;
    verifierOr
      : VerifierType -> VerifierType -> ComposeHashable VerifierType MapCon;
  };

projectVerifier
  {MapCon : Type -> Type}
  {OrdKey VerifierType Signable Commitment SignerType VerifierHashOrdKeyType : Type}
  (tc : ThresholdCompose
    OrdKey
    MapCon
    VerifierType
    Signable
    Commitment
    SignerType
    VerifierHashOrdKeyType)
  : Verifier
    VerifierHashOrdKeyType
    (ComposeHashable VerifierType MapCon)
    Signable
    (MapCon Commitment) :=
  mkVerifier@{
    verify := ThresholdCompose.verify tc;
    verifierHash := ThresholdCompose.verifierHash tc;
  };

ThresholdComposeFunctor
  {MapCon : Type -> Type}
  {OrdKey VerifierType Signable Commitment SignerType VerifierHashOrdKeyType : Type}
  (verifier : Verifier OrdKey VerifierType Signable Commitment)
  (signer : Signer SignerType Signable Commitment)
  (mapIn : OrdMap OrdKey MapCon)
  (thresholdComposeHash : HASH
    VerifierHashOrdKeyType
    (ComposeHashable VerifierType MapCon))
  : ThresholdCompose
    OrdKey
    MapCon
    VerifierType
    Signable
    Commitment
    SignerType
    VerifierHashOrdKeyType :=
  mkThresholdCompose@{
    map := mapIn;
    underlyingVerifier := verifier;
    underlyingSigner := signer;
    verifierHash := thresholdComposeHash;
    sign := \{s m := OrdMap.map map \{i := Signer.sign underlyingSigner i m} s};
    verify :=
      \{| (mkComposeHashable t ws) s c :=
        t
          <= OrdMap.foldl
            map
            \{(mkPair x y) := x + y}
            0
            (OrdMap.intersectWith
              map
              \{| (mkPair (mkPair w v) x) :=
                if 
                  | Verifier.verify underlyingVerifier v s x := w
                  | else := 0}
              (mkPair ws c))};
    
    signerCompose :=
      \{l :=
        foldl
          \{m (mkPair v s) :=
            OrdMap.insert
              map
              (mkPair
                m
                (mkPair
                  (HASH.hash (Verifier.verifierHash underlyingVerifier) v)
                  s))}
          (OrdMap.empty map)
          l};
    
    verifierCompose :=
      \{threshold weights :=
        mkComposeHashable
          threshold
          (foldl
            \{m (mkPair w v) :=
              OrdMap.insert
                map
                (mkPair
                  m
                  (mkPair
                    (HASH.hash (Verifier.verifierHash underlyingVerifier) v)
                    (mkPair w v)))}
            (OrdMap.empty map)
            weights)};
    
    verifierAnd := \{x y := verifierCompose 2 [mkPair 1 x; mkPair 1 y]};
    verifierOr := \{x y := verifierCompose 1 [mkPair 1 x; mkPair 1 y]};
  };

type ThresholdComposeSignsFor (OrdKey VerifierType Signable Commitment Evidence : Type) (MapCon : Type
  -> Type) VerifierHashOrdKeyType :=
  mkThresholdComposeSignsFor@{
    underlyingSignsFor
      : SignsFor OrdKey VerifierType Signable Commitment Evidence;
    verifier
      : ThresholdCompose
        OrdKey
        MapCon
        VerifierType
        Signable
        Commitment
        VerifierType
        VerifierHashOrdKeyType;
    signsFor
      : Evidence
        -> Pair
          (ComposeHashable VerifierType MapCon)
          (ComposeHashable VerifierType MapCon)
        -> Bool;
  };

projectSignsFor
  {OrdKey VerifierType Signable Commitment Evidence}
  {MapCon : Type -> Type}
  {VerifierHashOrdKeyType : Type}
  (tc : ThresholdComposeSignsFor
    OrdKey
    VerifierType
    Signable
    Commitment
    Evidence
    MapCon
    VerifierHashOrdKeyType)
  : SignsFor
    VerifierHashOrdKeyType
    (ComposeHashable VerifierType MapCon)
    Signable
    (MapCon Commitment)
    Evidence :=
  mkSignsFor@{
    verifier := projectVerifier (ThresholdComposeSignsFor.verifier tc);
    signsFor := ThresholdComposeSignsFor.signsFor tc;
  };

ThresholdComposeSignsForFunctor
  {OrdKey VerifierType Signable Commitment Evidence}
  {MapCon : Type -> Type}
  {VerifierHashOrdKeyType : Type}
  (S : SignsFor OrdKey VerifierType Signable Commitment Evidence)
  (signer : Signer VerifierType Signable Commitment)
  (map : OrdMap OrdKey MapCon)
  (thresholdComposeHash : HASH
    VerifierHashOrdKeyType
    (ComposeHashable VerifierType MapCon))
  : ThresholdComposeSignsFor
    OrdKey
    VerifierType
    Signable
    Commitment
    Evidence
    MapCon
    VerifierHashOrdKeyType :=
  mkThresholdComposeSignsFor@{
    underlyingSignsFor := S;
    verifier :=
      ThresholdComposeFunctor
        (SignsFor.verifier underlyingSignsFor)
        signer
        map
        thresholdComposeHash;
    signsFor :=
      \{e (mkPair (mkComposeHashable t0 w0) (mkComposeHashable t1 w1)) :=
        OrdMap.all
          map
          \{(mkPair w v) :=
            w * t1
              <= OrdMap.foldl
                  map
                  \{(mkPair (mkPair x v1) s) :=
                    if 
                      | SignsFor.signsFor underlyingSignsFor e (mkPair v v1) :=
                        x + s
                      | else := s}
                  0
                  w1
                * t0}
          w0};
  };

type ThresholdComposeEncryptor (OrdKey EncryptorType Plaintext Ciphertext : Type) (MapCon : Type
  -> Type) (EncryptorHashOrdKeyType : Type) :=
  mkThresholdComposeEncryptor@{
    map : OrdMap OrdKey MapCon;
    underlyingEncryptor : Encryptor OrdKey EncryptorType Plaintext Ciphertext;
    encryptorHash
      : HASH EncryptorHashOrdKeyType (ComposeHashable EncryptorType MapCon);
    compose
      : Nat
        -> List (Pair Nat EncryptorType)
        -> ComposeHashable EncryptorType MapCon;
    encrypt : ComposeHashable EncryptorType MapCon -> Plaintext -> Ciphertext;
  };

projectEncryptor
  {OrdKey EncryptorType Plaintext Ciphertext}
  {MapCon : Type -> Type}
  {EncryptorHashOrdKeyType}
  (tc : ThresholdComposeEncryptor
    OrdKey
    EncryptorType
    Plaintext
    Ciphertext
    MapCon
    EncryptorHashOrdKeyType)
  : Encryptor
    EncryptorHashOrdKeyType
    (ComposeHashable EncryptorType MapCon)
    Plaintext
    Ciphertext :=
  mkEncryptor@{
    encrypt := ThresholdComposeEncryptor.encrypt tc;
    encryptorHash := ThresholdComposeEncryptor.encryptorHash tc;
  };

axiom encrypt_DUMMY
  {EncryptorType Plaintext Ciphertext}
  {MapCon}
  : ComposeHashable EncryptorType MapCon -> Plaintext -> Ciphertext;

ThresholdComposeEncryptorFunctor
  {OrdKey EncryptorType Plaintext Ciphertext}
  {MapCon : Type -> Type}
  {EncryptorHashOrdKeyType}
  (encryptor : Encryptor OrdKey EncryptorType Plaintext Ciphertext)
  (mapIn : OrdMap OrdKey MapCon)
  (thresholdComposeHash : HASH
    EncryptorHashOrdKeyType
    (ComposeHashable EncryptorType MapCon))
  : ThresholdComposeEncryptor
    OrdKey
    EncryptorType
    Plaintext
    Ciphertext
    MapCon
    EncryptorHashOrdKeyType :=
  mkThresholdComposeEncryptor@{
    map := mapIn;
    underlyingEncryptor := encryptor;
    encryptorHash := thresholdComposeHash;
    compose :=
      \{t w :=
        mkComposeHashable@{
          threshold := t;
          weights :=
            foldl
              \{m (mkPair w e) :=
                OrdMap.insert
                  map
                  (mkPair
                    m
                    (mkPair
                      (HASH.hash
                        (Encryptor.encryptorHash underlyingEncryptor)
                        e)
                      (mkPair w e)))}
              (OrdMap.empty map)
              w;
        }};
    encrypt := encrypt_DUMMY;
  };

type ThresholdComposeReadsFor (OrdKey EncryptorType Plaintext Ciphertext Evidence : Type) (MapCon : Type
  -> Type) (EncryptorHashOrdKeyType : Type) :=
  mkThresholdComposeReadsFor@{
    underlyingReadsFor
      : ReadsFor OrdKey EncryptorType Plaintext Ciphertext Evidence;
    encryptor
      : ThresholdComposeEncryptor
        OrdKey
        EncryptorType
        Plaintext
        Ciphertext
        MapCon
        EncryptorHashOrdKeyType;
    readsFor
      : Evidence
        -> Pair
          (ComposeHashable EncryptorType MapCon)
          (ComposeHashable EncryptorType MapCon)
        -> Bool;
  };

projectReadsFor
  {OrdKey VerifierType Signable Commitment Evidence}
  {MapCon : Type -> Type}
  {EncryptorHashOrdKeyType : Type}
  (tc : ThresholdComposeReadsFor
    OrdKey
    VerifierType
    Signable
    Commitment
    Evidence
    MapCon
    EncryptorHashOrdKeyType)
  : ReadsFor
    EncryptorHashOrdKeyType
    (ComposeHashable VerifierType MapCon)
    Signable
    Commitment
    Evidence :=
  mkReadsFor@{
    encryptor := projectEncryptor (ThresholdComposeReadsFor.encryptor tc);
    readsFor := ThresholdComposeReadsFor.readsFor tc;
  };

ThresholdComposeReadsForFunctor
  {OrdKey EncryptorType Plaintext Ciphertext Evidence}
  {MapCon : Type -> Type}
  {EncryptorHashOrdKeyType}
  (r : ReadsFor OrdKey EncryptorType Plaintext Ciphertext Evidence)
  (map : OrdMap OrdKey MapCon)
  (thresholdComposeHash : HASH
    EncryptorHashOrdKeyType
    (ComposeHashable EncryptorType MapCon))
  : ThresholdComposeReadsFor
    OrdKey
    EncryptorType
    Plaintext
    Ciphertext
    Evidence
    MapCon
    EncryptorHashOrdKeyType :=
  mkThresholdComposeReadsFor@{
    underlyingReadsFor := r;
    encryptor :=
      ThresholdComposeEncryptorFunctor
        (ReadsFor.encryptor underlyingReadsFor)
        map
        thresholdComposeHash;
    readsFor :=
      \{e (mkPair (mkComposeHashable t0 w0) (mkComposeHashable t1 w1)) :=
        OrdMap.all
          map
          \{(mkPair w v) :=
            w * t1
              <= OrdMap.foldl
                  map
                  \{(mkPair (mkPair x v1) s) :=
                    if 
                      | ReadsFor.readsFor underlyingReadsFor e (mkPair v v1) :=
                        x + s
                      | else := s}
                  0
                  w1
                * t0}
          w0};
  };

type VerifierName OrdKey VerifierType Signable Commitment Evidence IdentityName VerifierNameHashOrdKeyType :=
  mkVerifierName@{
    verifier : Verifier OrdKey VerifierType Signable Commitment;
    checkVerifierName : IdentityName -> VerifierType -> Evidence -> Bool;
    verifierNameHash : HASH VerifierNameHashOrdKeyType IdentityName;
  };

type EncryptorName OrdKey EncryptorType Plaintext Ciphertext Evidence IdentityName EncryptorNameHashOrdKeyType :=
  mkEncryptorName@{
    verifier : Encryptor OrdKey EncryptorType Plaintext Ciphertext;
    checkEncryptorName : IdentityName -> EncryptorType -> Evidence -> Bool;
    encryptorNameHash : HASH EncryptorNameHashOrdKeyType IdentityName;
  };

SubVerifierFunctor
  (OrdKey VerifierType Signable Commitment Name ParentOrdKeyType : Type)
  (child : Verifier OrdKey VerifierType Signable Commitment)
  (parent : Verifier
    ParentOrdKeyType
    VerifierType
    (Pair String (Pair Name OrdKey))
    Commitment)
  (hash : HASH ParentOrdKeyType (Pair ParentOrdKeyType Name))
  : VerifierName
    OrdKey
    VerifierType
    Signable
    Commitment
    (Pair VerifierType Commitment)
    (Pair ParentOrdKeyType Name)
    ParentOrdKeyType :=
  mkVerifierName@{
    verifier := child;
    checkVerifierName :=
      \{(mkPair ph n) c (mkPair pv pc) :=
        Verifier.verify
            parent
            pv
            (mkPair
              "I identify this verifier with this name : "
              (mkPair n (HASH.hash (Verifier.verifierHash child) c)))
            pc
          && OrdKey.compare
              (HASH.ordKey (Verifier.verifierHash parent))
              ph
              (HASH.hash (Verifier.verifierHash parent) pv)
            == Equal};
    verifierNameHash := hash;
  };