module arch.node.engines.identity_management_behaviour;

import prelude open;
import arch.node.engines.commitment_config open;
import arch.node.engines.decryption_config open;
import arch.node.engines.commitment_environment open;
import arch.node.engines.decryption_environment open;
import arch.node.engines.identity_management_environment open;
import arch.node.engines.identity_management_messages open;
import arch.node.engines.identity_management_config open;
import arch.node.types.anoma as Anoma open;
import arch.node.types.engine open;
import arch.node.types.engine_environment open;
import arch.node.types.identities open;
import arch.node.types.messages open;
import arch.system.identity.identity open hiding {ExternalIdentity};

type MessageFrom :=
  mkMessageFrom@{
    whoAsked : Option EngineID;
    mailbox : Option MailboxID;
  };

type IdentityManagementActionArgument :=
  | IdentityManagementActionArgumentMessageFrom MessageFrom;

IdentityManagementActionArguments : Type :=
  List IdentityManagementActionArgument;

IdentityManagementAction : Type :=
  Action
    IdentityManagementCfg
    IdentityManagementLocalState
    IdentityManagementMailboxState
    IdentityManagementTimerHandle
    IdentityManagementActionArguments
    Anoma.Msg
    Anoma.Cfg
    Anoma.Env;

IdentityManagementActionInput : Type :=
  ActionInput
    IdentityManagementCfg
    IdentityManagementLocalState
    IdentityManagementMailboxState
    IdentityManagementTimerHandle
    IdentityManagementActionArguments
    Anoma.Msg;

IdentityManagementActionEffect : Type :=
  ActionEffect
    IdentityManagementLocalState
    IdentityManagementMailboxState
    IdentityManagementTimerHandle
    Anoma.Msg
    Anoma.Cfg
    Anoma.Env;

IdentityManagementActionExec : Type :=
  ActionExec
    IdentityManagementCfg
    IdentityManagementLocalState
    IdentityManagementMailboxState
    IdentityManagementTimerHandle
    IdentityManagementActionArguments
    Anoma.Msg
    Anoma.Cfg
    Anoma.Env;

hasCommitCapability (capabilities : Capabilities) : Bool :=
  case capabilities of
    | CapabilityCommit := true
    | CapabilityCommitAndDecrypt := true
    | _ := false;

hasDecryptCapability (capabilities : Capabilities) : Bool :=
  case capabilities of
    | CapabilityDecrypt := true
    | CapabilityCommitAndDecrypt := true
    | _ := false;

isSubsetCapabilities
  (requested : Capabilities) (available : Capabilities) : Bool :=
  (not (hasCommitCapability requested) || hasCommitCapability available)
    && not (hasDecryptCapability requested)
    || hasDecryptCapability available;

updateIdentityAndSpawnEngines
  (env : IdentityManagementEnv)
  (backend' : Backend)
  (whoAsked : EngineID)
  (identityInfo : IdentityInfo)
  (capabilities' : Capabilities)
  : Pair IdentityInfo (List (Pair Cfg Env)) :=
  let
    decryptionConfig : DecryptionCfg :=
      mkDecryptionCfg@{
        decryptor := genDecryptor backend';
        backend := backend';
      };
    decryptionEnv : DecryptionEnv :=
      mkEngineEnv@{
        localState := unit;
        mailboxCluster := Map.empty;
        acquaintances := Set.empty;
        timers := [];
      };
    decryptionEng : Pair Cfg Env :=
      mkPair (CfgDecryption decryptionConfig) (EnvDecryption decryptionEnv);
    commitmentConfig : CommitmentCfg :=
      mkCommitmentCfg@{
        signer := genSigner backend';
        backend := backend';
      };
    commitmentEnv : CommitmentEnv :=
      mkEngineEnv@{
        localState := unit;
        mailboxCluster := Map.empty;
        acquaintances := Set.empty;
        timers := [];
      };
    commitmentEng : Pair Cfg Env :=
      mkPair (CfgCommitment commitmentConfig) (EnvCommitment commitmentEnv);
  in case capabilities' of
       | CapabilityCommitAndDecrypt :=
         let
           spawnedEngines := [decryptionEng; commitmentEng];
           commitmentEngineName := nameGen "committer" (snd whoAsked) whoAsked;
           decryptionEngineName := nameGen "decryptor" (snd whoAsked) whoAsked;
           updatedIdentityInfo1 :=
             identityInfo@IdentityInfo{
               commitmentEngine := some (mkPair none commitmentEngineName);
               decryptionEngine := some (mkPair none decryptionEngineName);
             };
         in mkPair updatedIdentityInfo1 spawnedEngines
       | CapabilityCommit :=
         let
           spawnedEngines := [commitmentEng];
           commitmentEngineName := nameGen "committer" (snd whoAsked) whoAsked;
           updatedIdentityInfo1 :=
             identityInfo@IdentityInfo{commitmentEngine := some
               (mkPair none commitmentEngineName)};
         in mkPair updatedIdentityInfo1 spawnedEngines
       | CapabilityDecrypt :=
         let
           spawnedEngines := [decryptionEng];
           decryptionEngineName := nameGen "decryptor" (snd whoAsked) whoAsked;
           updatedIdentityInfo1 :=
             identityInfo@IdentityInfo{decryptionEngine := some
               (mkPair none decryptionEngineName)};
         in mkPair updatedIdentityInfo1 spawnedEngines;

copyEnginesForCapabilities
  (env : IdentityManagementEnv)
  (whoAsked : EngineID)
  (externalIdentityInfo : IdentityInfo)
  (requestedCapabilities : Capabilities)
  : IdentityInfo :=
  let
    newIdentityInfo :=
      mkIdentityInfo@{
        backend := IdentityInfo.backend externalIdentityInfo;
        capabilities := requestedCapabilities;
        commitmentEngine :=
          case hasCommitCapability requestedCapabilities of
            | true := IdentityInfo.commitmentEngine externalIdentityInfo
            | false := none;
        decryptionEngine :=
          case hasDecryptCapability requestedCapabilities of
            | true := IdentityInfo.decryptionEngine externalIdentityInfo
            | false := none;
      };
  in newIdentityInfo;

generateIdentityAction
  (input : IdentityManagementActionInput)
  : Option IdentityManagementActionEffect :=
  let
    env := ActionInput.env input;
    local := EngineEnv.localState env;
    identities := IdentityManagementLocalState.identities local;
    trigger := ActionInput.trigger input;
  in case getEngineMsgFromTimestampedTrigger trigger of
       | some emsg :=
         let
           whoAsked := EngineMsg.sender emsg;
         in case Map.lookup whoAsked identities of {
              | some _ :=
                some
                  mkActionEffect@{
                    env := env;
                    msgs :=
                      [
                        mkEngineMsg@{
                          sender :=
                            getEngineIDFromEngineCfg (ActionInput.cfg input);
                          target := whoAsked;
                          mailbox := some 0;
                          msg :=
                            MsgIdentityManagement
                              (MsgIdentityManagementGenerateIdentityReply
                                mkReplyGenerateIdentity@{
                                  commitmentEngine := none;
                                  decryptionEngine := none;
                                  externalIdentity := whoAsked;
                                  err := some "Identity already exists";
                                });
                        };
                      ];
                    timers := [];
                    engines := [];
                  }
              | none :=
                case emsg of {
                  | mkEngineMsg@{
                      msg := Anoma.MsgIdentityManagement (MsgIdentityManagementGenerateIdentityRequest (mkRequestGenerateIdentity backend' params' capabilities'));
                    } :=
                    let
                      identityInfo :=
                        mkIdentityInfo@{
                          backend := backend';
                          capabilities := capabilities';
                          commitmentEngine := none;
                          decryptionEngine := none;
                        };
                      pair' :=
                        updateIdentityAndSpawnEngines
                          env
                          backend'
                          whoAsked
                          identityInfo
                          capabilities';
                      updatedIdentityInfo := fst pair';
                      spawnedEnginesFinal := snd pair';
                      updatedIdentities :=
                        Map.insert whoAsked updatedIdentityInfo identities;
                      newLocalState :=
                        local@IdentityManagementLocalState{identities := updatedIdentities};
                      newEnv' := env@EngineEnv{localState := newLocalState};
                    in some
                      mkActionEffect@{
                        env := newEnv';
                        msgs :=
                          [
                            mkEngineMsg@{
                              sender :=
                                getEngineIDFromEngineCfg
                                  (ActionInput.cfg input);
                              target := whoAsked;
                              mailbox := some 0;
                              msg :=
                                MsgIdentityManagement
                                  (MsgIdentityManagementGenerateIdentityReply
                                    mkReplyGenerateIdentity@{
                                      commitmentEngine :=
                                        IdentityInfo.commitmentEngine
                                          updatedIdentityInfo;
                                      decryptionEngine :=
                                        IdentityInfo.decryptionEngine
                                          updatedIdentityInfo;
                                      externalIdentity := whoAsked;
                                      err := none;
                                    });
                            };
                          ];
                        timers := [];
                        engines := spawnedEnginesFinal;
                      }
                  | _ := none
                }
            }
       | _ := none;

connectIdentityAction
  (input : IdentityManagementActionInput)
  : Option IdentityManagementActionEffect :=
  let
    env := ActionInput.env input;
    local := EngineEnv.localState env;
    identities := IdentityManagementLocalState.identities local;
    trigger := ActionInput.trigger input;
  in case getEngineMsgFromTimestampedTrigger trigger of
       | some emsg :=
         let
           whoAsked := EngineMsg.sender emsg;
         in case Map.lookup whoAsked identities of {
              | some _ :=
                some
                  mkActionEffect@{
                    env := env;
                    msgs :=
                      [
                        mkEngineMsg@{
                          sender :=
                            getEngineIDFromEngineCfg (ActionInput.cfg input);
                          target := whoAsked;
                          mailbox := some 0;
                          msg :=
                            MsgIdentityManagement
                              (MsgIdentityManagementConnectIdentityReply
                                mkConnectIdentityReply@{
                                  commitmentEngine := none;
                                  decryptionEngine := none;
                                  err := some "Identity already exists";
                                });
                        };
                      ];
                    timers := [];
                    engines := [];
                  }
              | none :=
                case emsg of {
                  | mkEngineMsg@{
                      msg := Anoma.MsgIdentityManagement (MsgIdentityManagementConnectIdentityRequest (mkRequestConnectIdentity externalIdentity' backend' capabilities'));
                    } :=
                    case Map.lookup externalIdentity' identities of {
                      | none :=
                        some
                          mkActionEffect@{
                            env := env;
                            msgs :=
                              [
                                mkEngineMsg@{
                                  sender :=
                                    getEngineIDFromEngineCfg
                                      (ActionInput.cfg input);
                                  target := whoAsked;
                                  mailbox := some 0;
                                  msg :=
                                    MsgIdentityManagement
                                      (MsgIdentityManagementConnectIdentityReply
                                        mkConnectIdentityReply@{
                                          commitmentEngine := none;
                                          decryptionEngine := none;
                                          err :=
                                            some "External identity not found";
                                        });
                                };
                              ];
                            timers := [];
                            engines := [];
                          }
                      | some externalIdentityInfo :=
                        if 
                          | isSubsetCapabilities
                            capabilities'
                            (IdentityInfo.capabilities externalIdentityInfo) :=
                            let
                              newIdentityInfo :=
                                copyEnginesForCapabilities
                                  env
                                  whoAsked
                                  externalIdentityInfo
                                  capabilities';
                              updatedIdentities :=
                                Map.insert whoAsked newIdentityInfo identities;
                              newLocalState :=
                                local@IdentityManagementLocalState{identities := updatedIdentities};
                              newEnv' :=
                                env@EngineEnv{localState := newLocalState};
                            in some
                              mkActionEffect@{
                                env := newEnv';
                                msgs :=
                                  [
                                    mkEngineMsg@{
                                      sender :=
                                        getEngineIDFromEngineCfg
                                          (ActionInput.cfg input);
                                      target := whoAsked;
                                      mailbox := some 0;
                                      msg :=
                                        MsgIdentityManagement
                                          (MsgIdentityManagementConnectIdentityReply
                                            mkConnectIdentityReply@{
                                              commitmentEngine :=
                                                IdentityInfo.commitmentEngine
                                                  newIdentityInfo;
                                              decryptionEngine :=
                                                IdentityInfo.decryptionEngine
                                                  newIdentityInfo;
                                              err := none;
                                            });
                                    };
                                  ];
                                timers := [];
                                engines := [];
                              }
                          | else :=
                            some
                              mkActionEffect@{
                                env := env;
                                msgs :=
                                  [
                                    mkEngineMsg@{
                                      sender :=
                                        getEngineIDFromEngineCfg
                                          (ActionInput.cfg input);
                                      target := whoAsked;
                                      mailbox := some 0;
                                      msg :=
                                        MsgIdentityManagement
                                          (MsgIdentityManagementConnectIdentityReply
                                            mkConnectIdentityReply@{
                                              commitmentEngine := none;
                                              decryptionEngine := none;
                                              err :=
                                                some
                                                  "Capabilities not available";
                                            });
                                    };
                                  ];
                                timers := [];
                                engines := [];
                              }
                    }
                  | _ := none
                }
            }
       | _ := none;

deleteIdentityAction
  (input : IdentityManagementActionInput)
  : Option IdentityManagementActionEffect :=
  let
    env := ActionInput.env input;
    local := EngineEnv.localState env;
    identities := IdentityManagementLocalState.identities local;
    trigger := ActionInput.trigger input;
  in case getEngineMsgFromTimestampedTrigger trigger of
       | some emsg :=
         let
           whoAsked := EngineMsg.sender emsg;
         in case emsg of {
              | mkEngineMsg@{
                  msg := Anoma.MsgIdentityManagement (MsgIdentityManagementDeleteIdentityRequest (mkRequestDeleteIdentity externalIdentity backend'));
                } :=
                case Map.lookup externalIdentity identities of {
                  | none :=
                    some
                      mkActionEffect@{
                        env := env;
                        msgs :=
                          [
                            mkEngineMsg@{
                              sender :=
                                getEngineIDFromEngineCfg
                                  (ActionInput.cfg input);
                              target := whoAsked;
                              mailbox := some 0;
                              msg :=
                                MsgIdentityManagement
                                  (MsgIdentityManagementDeleteIdentityReply
                                    mkReplyDeleteIdentity@{
                                      err := some "Identity does not exist";
                                    });
                            };
                          ];
                        timers := [];
                        engines := [];
                      }
                  | some _ :=
                    let
                      updatedIdentities :=
                        Map.delete externalIdentity identities;
                      newLocalState :=
                        local@IdentityManagementLocalState{identities := updatedIdentities};
                      newEnv' := env@EngineEnv{localState := newLocalState};
                    in some
                      mkActionEffect@{
                        env := newEnv';
                        msgs :=
                          [
                            mkEngineMsg@{
                              sender :=
                                getEngineIDFromEngineCfg
                                  (ActionInput.cfg input);
                              target := whoAsked;
                              mailbox := some 0;
                              msg :=
                                MsgIdentityManagement
                                  (MsgIdentityManagementDeleteIdentityReply
                                    mkReplyDeleteIdentity@{
                                      err := none;
                                    });
                            };
                          ];
                        timers := [];
                        engines := [];
                      }
                }
              | _ := none
            }
       | _ := none;

generateIdentityActionLabel : IdentityManagementActionExec :=
  Seq [generateIdentityAction];

connectIdentityActionLabel : IdentityManagementActionExec :=
  Seq [connectIdentityAction];

deleteIdentityActionLabel : IdentityManagementActionExec :=
  Seq [deleteIdentityAction];

IdentityManagementGuard : Type :=
  Guard
    IdentityManagementCfg
    IdentityManagementLocalState
    IdentityManagementMailboxState
    IdentityManagementTimerHandle
    IdentityManagementActionArguments
    Anoma.Msg
    Anoma.Cfg
    Anoma.Env;

IdentityManagementGuardOutput : Type :=
  GuardOutput
    IdentityManagementCfg
    IdentityManagementLocalState
    IdentityManagementMailboxState
    IdentityManagementTimerHandle
    IdentityManagementActionArguments
    Anoma.Msg
    Anoma.Cfg
    Anoma.Env;

IdentityManagementGuardEval : Type :=
  GuardEval
    IdentityManagementCfg
    IdentityManagementLocalState
    IdentityManagementMailboxState
    IdentityManagementTimerHandle
    IdentityManagementActionArguments
    Anoma.Msg
    Anoma.Cfg
    Anoma.Env;

generateIdentityGuard
  (trigger : TimestampedTrigger IdentityManagementTimerHandle Anoma.Msg)
  (cfg : EngineCfg IdentityManagementCfg)
  (env : IdentityManagementEnv)
  : Option IdentityManagementGuardOutput :=
  case getEngineMsgFromTimestampedTrigger trigger of
    | some mkEngineMsg@{
             msg := Anoma.MsgIdentityManagement (MsgIdentityManagementGenerateIdentityRequest _);
           } :=
      some
        mkGuardOutput@{
          action := generateIdentityActionLabel;
          args := [];
        }
    | _ := none;

connectIdentityGuard
  (trigger : TimestampedTrigger IdentityManagementTimerHandle Anoma.Msg)
  (cfg : EngineCfg IdentityManagementCfg)
  (env : IdentityManagementEnv)
  : Option IdentityManagementGuardOutput :=
  case getEngineMsgFromTimestampedTrigger trigger of
    | some mkEngineMsg@{
             msg := Anoma.MsgIdentityManagement (MsgIdentityManagementConnectIdentityRequest _);
           } :=
      some
        mkGuardOutput@{
          action := connectIdentityActionLabel;
          args := [];
        }
    | _ := none;

deleteIdentityGuard
  (trigger : TimestampedTrigger IdentityManagementTimerHandle Anoma.Msg)
  (cfg : EngineCfg IdentityManagementCfg)
  (env : IdentityManagementEnv)
  : Option IdentityManagementGuardOutput :=
  case getEngineMsgFromTimestampedTrigger trigger of
    | some mkEngineMsg@{
             msg := Anoma.MsgIdentityManagement (MsgIdentityManagementDeleteIdentityRequest _);
           } :=
      some
        mkGuardOutput@{
          action := deleteIdentityActionLabel;
          args := [];
        }
    | _ := none;

IdentityManagementBehaviour : Type :=
  EngineBehaviour
    IdentityManagementCfg
    IdentityManagementLocalState
    IdentityManagementMailboxState
    IdentityManagementTimerHandle
    IdentityManagementActionArguments
    Anoma.Msg
    Anoma.Cfg
    Anoma.Env;

identityManagementBehaviour : IdentityManagementBehaviour :=
  mkEngineBehaviour@{
    guards :=
      First [generateIdentityGuard; connectIdentityGuard; deleteIdentityGuard];
  };