module arch.node.engines.encryption_behaviour;

import prelude open;
import arch.system.identity.identity open hiding {ExternalIdentity};
import arch.node.engines.encryption_environment open;
import arch.node.engines.encryption_messages open;
import arch.node.engines.encryption_config open;
import arch.node.engines.reads_for_messages open;
import arch.node.types.anoma as Anoma open;
import arch.node.types.engine open;
import arch.node.types.identities open;
import arch.node.types.messages open;

type ReplyTo :=
  mkReplyTo@{
    whoAsked : Option EngineID;
    mailbox : Option MailboxID;
  };

type EncryptionActionArgument := | EncryptionActionArgumentReplyTo ReplyTo;

EncryptionActionArguments : Type := List EncryptionActionArgument;

EncryptionAction : Type :=
  Action
    EncryptionCfg
    EncryptionLocalState
    EncryptionMailboxState
    EncryptionTimerHandle
    EncryptionActionArguments
    Anoma.Msg
    Anoma.Cfg
    Anoma.Env;

EncryptionActionInput : Type :=
  ActionInput
    EncryptionCfg
    EncryptionLocalState
    EncryptionMailboxState
    EncryptionTimerHandle
    EncryptionActionArguments
    Anoma.Msg;

EncryptionActionEffect : Type :=
  ActionEffect
    EncryptionLocalState
    EncryptionMailboxState
    EncryptionTimerHandle
    Anoma.Msg
    Anoma.Cfg
    Anoma.Env;

EncryptionActionExec : Type :=
  ActionExec
    EncryptionCfg
    EncryptionLocalState
    EncryptionMailboxState
    EncryptionTimerHandle
    EncryptionActionArguments
    Anoma.Msg
    Anoma.Cfg
    Anoma.Env;

encryptAction (input : EncryptionActionInput) : Option EncryptionActionEffect :=
  let
    env := ActionInput.env input;
    cfg := ActionInput.cfg input;
    tt := ActionInput.trigger input;
    localState := EngineEnv.localState env;
  in case getEngineMsgFromTimestampedTrigger tt of
       | some emsg :=
         case EngineMsg.msg emsg of {
           | Anoma.MsgEncryption (MsgEncryptionRequest (mkRequestEncrypt data externalIdentity useReadsFor)) :=
             case useReadsFor of {
               | false :=
                 some
                   mkActionEffect@{
                     env := env;
                     msgs :=
                       [
                         mkEngineMsg@{
                           sender := getEngineIDFromEngineCfg cfg;
                           target := EngineMsg.sender emsg;
                           mailbox := some 0;
                           msg :=
                             Anoma.MsgEncryption
                               (MsgEncryptionReply
                                 mkReplyEncrypt@{
                                   ciphertext :=
                                     Encryptor.encrypt
                                       (EncryptionCfg.encryptor
                                         (EngineCfg.cfg cfg)
                                         Set.empty
                                         externalIdentity)
                                       (EncryptionCfg.backend
                                         (EngineCfg.cfg cfg))
                                       data;
                                   err := none;
                                 });
                         };
                       ];
                     timers := [];
                     engines := [];
                   }
               | true :=
                 let
                   existingRequests :=
                     Map.lookup
                       externalIdentity
                       (EncryptionLocalState.pendingRequests localState);
                   newPendingList :=
                     case existingRequests of
                       | some reqs :=
                         reqs ++ [mkPair (EngineMsg.sender emsg) data]
                       | none := [mkPair (EngineMsg.sender emsg) data];
                   newLocalState :=
                     localState@EncryptionLocalState{pendingRequests := Map.insert
                       externalIdentity
                       newPendingList
                       (EncryptionLocalState.pendingRequests localState)};
                 in some
                   mkActionEffect@{
                     env := env@EngineEnv{localState := newLocalState};
                     msgs :=
                       case existingRequests of
                         | some _ := []
                         | none :=
                           [
                             mkEngineMsg@{
                               sender := getEngineIDFromEngineCfg cfg;
                               target :=
                                 EncryptionCfg.readsForEngineAddress
                                   (EngineCfg.cfg cfg);
                               mailbox := some 0;
                               msg :=
                                 Anoma.MsgReadsFor
                                   (MsgQueryReadsForEvidenceRequest
                                     mkRequestQueryReadsForEvidence@{
                                       externalIdentity := externalIdentity;
                                     });
                             };
                           ];
                     timers := [];
                     engines := [];
                   }
             }
           | _ := none
         }
       | _ := none;

handleReadsForReplyAction
  (input : EncryptionActionInput) : Option EncryptionActionEffect :=
  let
    env := ActionInput.env input;
    cfg := ActionInput.cfg input;
    tt := ActionInput.trigger input;
    localState := EngineEnv.localState env;
  in case getEngineMsgFromTimestampedTrigger tt of
       | some emsg :=
         case EngineMsg.msg emsg of {
           | Anoma.MsgReadsFor (MsgQueryReadsForEvidenceReply (mkReplyQueryReadsForEvidence externalIdentity evidence err)) :=
             case
               Map.lookup
                 externalIdentity
                 (EncryptionLocalState.pendingRequests localState)
             of {
               | some reqs :=
                 let
                   newLocalState :=
                     localState@EncryptionLocalState{pendingRequests := Map.delete
                       externalIdentity
                       (EncryptionLocalState.pendingRequests localState)};
                 in some
                   mkActionEffect@{
                     env := env@EngineEnv{localState := newLocalState};
                     msgs :=
                       map
                         \{req :=
                           let
                             whoAsked := fst req;
                             data := snd req;
                           in mkEngineMsg@{
                                sender := getEngineIDFromEngineCfg cfg;
                                target := whoAsked;
                                mailbox := some 0;
                                msg :=
                                  Anoma.MsgEncryption
                                    (MsgEncryptionReply
                                      mkReplyEncrypt@{
                                        ciphertext :=
                                          Encryptor.encrypt
                                            (EncryptionCfg.encryptor
                                              (EngineCfg.cfg cfg)
                                              evidence
                                              externalIdentity)
                                            (EncryptionCfg.backend
                                              (EngineCfg.cfg cfg))
                                            data;
                                        err := none;
                                      });
                              }}
                         reqs;
                     timers := [];
                     engines := [];
                   }
               | none := none
             }
           | _ := none
         }
       | _ := none;

encryptActionLabel : EncryptionActionExec := Seq [encryptAction];

handleReadsForReplyActionLabel : EncryptionActionExec :=
  Seq [handleReadsForReplyAction];

EncryptionGuard : Type :=
  Guard
    EncryptionCfg
    EncryptionLocalState
    EncryptionMailboxState
    EncryptionTimerHandle
    EncryptionActionArguments
    Anoma.Msg
    Anoma.Cfg
    Anoma.Env;

EncryptionGuardOutput : Type :=
  GuardOutput
    EncryptionCfg
    EncryptionLocalState
    EncryptionMailboxState
    EncryptionTimerHandle
    EncryptionActionArguments
    Anoma.Msg
    Anoma.Cfg
    Anoma.Env;

encryptGuard
  (tt : TimestampedTrigger EncryptionTimerHandle Anoma.Msg)
  (cfg : EngineCfg EncryptionCfg)
  (env : EncryptionEnv)
  : Option EncryptionGuardOutput :=
  case getEngineMsgFromTimestampedTrigger tt of
    | some mkEngineMsg@{msg := Anoma.MsgEncryption (MsgEncryptionRequest _)} :=
      some
        mkGuardOutput@{
          action := encryptActionLabel;
          args := [];
        }
    | _ := none;

readsForReplyGuard
  (tt : TimestampedTrigger EncryptionTimerHandle Anoma.Msg)
  (cfg : EngineCfg EncryptionCfg)
  (env : EncryptionEnv)
  : Option EncryptionGuardOutput :=
  case getEngineMsgFromTimestampedTrigger tt of
    | some emsg :=
      case EngineMsg.msg emsg of {
        | Anoma.MsgReadsFor (MsgQueryReadsForEvidenceReply _) :=
          case
            isEqual
              (Ord.cmp
                (EngineMsg.sender emsg)
                (EncryptionCfg.readsForEngineAddress (EngineCfg.cfg cfg)))
          of {
            | true :=
              some
                mkGuardOutput@{
                  action := handleReadsForReplyActionLabel;
                  args := [];
                }
            | false := none
          }
        | _ := none
      }
    | _ := none;

EncryptionBehaviour : Type :=
  EngineBehaviour
    EncryptionCfg
    EncryptionLocalState
    EncryptionMailboxState
    EncryptionTimerHandle
    EncryptionActionArguments
    Anoma.Msg
    Anoma.Cfg
    Anoma.Env;

encryptionBehaviour : EncryptionBehaviour :=
  mkEngineBehaviour@{
    guards := First [encryptGuard; readsForReplyGuard];
  };