/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

syntax = "proto3";

package org.apache.spark.sql.execution.streaming.state;

message StateRequest {
  int32 version = 1;
  oneof method {
    StatefulProcessorCall statefulProcessorCall = 2;
    StateVariableRequest stateVariableRequest = 3;
    ImplicitGroupingKeyRequest implicitGroupingKeyRequest = 4;
    TimerRequest timerRequest = 5;
    UtilsRequest utilsRequest = 6;
  }
}

message StateResponse {
  int32 statusCode = 1;
  string errorMessage = 2;
  bytes value = 3;
}

message StateResponseWithLongTypeVal {
  int32 statusCode = 1;
  string errorMessage = 2;
  int64 value = 3;
}

message StateResponseWithStringTypeVal {
  int32 statusCode = 1;
  string errorMessage = 2;
  string value = 3;
}

message StateResponseWithListGet {
  int32 statusCode = 1;
  string errorMessage = 2;
  repeated bytes value = 3;
  bool requireNextFetch = 4;
}

message StateResponseWithMapKeysOrValues {
  int32 statusCode = 1;
  string errorMessage = 2;
  repeated bytes value = 3;
  bool requireNextFetch = 4;
}

message KeyAndValuePair {
  bytes key = 1;
  bytes value = 2;
}

message StateResponseWithMapIterator {
  int32 statusCode = 1;
  string errorMessage = 2;
  repeated KeyAndValuePair kvPair = 3;
  bool requireNextFetch = 4;
}

message TimerInfo {
  optional bytes key = 1;
  int64 timestampMs = 2;
}

message StateResponseWithTimer {
  int32 statusCode = 1;
  string errorMessage = 2;
  repeated TimerInfo timer = 3;
  bool requireNextFetch = 4;
}

message StatefulProcessorCall {
  oneof method {
    SetHandleState setHandleState = 1;
    StateCallCommand getValueState = 2;
    StateCallCommand getListState = 3;
    StateCallCommand getMapState = 4;
    TimerStateCallCommand timerStateCall = 5;
    StateCallCommand deleteIfExists = 6;
  }
}

message StateVariableRequest {
  oneof method {
    ValueStateCall valueStateCall = 1;
    ListStateCall listStateCall = 2;
    MapStateCall mapStateCall = 3;
  }
}

message ImplicitGroupingKeyRequest {
  oneof method {
    SetImplicitKey setImplicitKey = 1;
    RemoveImplicitKey removeImplicitKey = 2;
  }
}

message TimerRequest {
  oneof method {
    TimerValueRequest timerValueRequest = 1;
    ExpiryTimerRequest expiryTimerRequest = 2;
  }
}

message TimerValueRequest {
  oneof method {
    GetProcessingTime getProcessingTimer = 1;
    GetWatermark getWatermark = 2;
  }
}

message ExpiryTimerRequest {
  string iteratorId = 1;
  int64 expiryTimestampMs = 2;
}

message GetProcessingTime {
}

message GetWatermark {
}

message UtilsRequest {
  oneof method {
    ParseStringSchema parseStringSchema = 1;
  }
}

message ParseStringSchema {
  string schema = 1;
}

message StateCallCommand {
  string stateName = 1;
  string schema = 2;
  string mapStateValueSchema = 3;
  TTLConfig ttl = 4;
}

message TimerStateCallCommand {
  oneof method {
    RegisterTimer register = 1;
    DeleteTimer delete = 2;
    ListTimers list = 3;
  }
}

message ValueStateCall {
  string stateName = 1;
  oneof method {
    Exists exists = 2;
    Get get = 3;
    ValueStateUpdate valueStateUpdate = 4;
    Clear clear = 5;
  }
}

message ListStateCall {
  string stateName = 1;
  oneof method {
    Exists exists = 2;
    ListStateGet listStateGet = 3;
    ListStatePut listStatePut = 4;
    AppendValue appendValue = 5;
    AppendList appendList = 6;
    Clear clear = 7;
  }
}

message MapStateCall {
  string stateName = 1;
  oneof method {
    Exists exists = 2;
    GetValue getValue = 3;
    ContainsKey containsKey = 4;
    UpdateValue updateValue = 5;
    Iterator iterator = 6;
    Keys keys = 7;
    Values values = 8;
    RemoveKey removeKey = 9;
    Clear clear = 10;
  }
}

message SetImplicitKey {
  bytes key = 1;
}

message RemoveImplicitKey {
}

message Exists {
}

message Get {
}

message RegisterTimer {
  int64 expiryTimestampMs = 1;
}

message DeleteTimer {
  int64 expiryTimestampMs = 1;
}

message ListTimers {
  string iteratorId = 1;
}

message ValueStateUpdate {
  bytes value = 1;
}

message Clear {
}

message ListStateGet {
  string iteratorId = 1;
}

message ListStatePut {
  repeated bytes value = 1;
  bool fetchWithArrow = 2;
}

message AppendValue {
  bytes value = 1;
}

message AppendList {
  repeated bytes value = 1;
  bool fetchWithArrow = 2;
}

message GetValue {
  bytes userKey = 1;
}

message ContainsKey {
  bytes userKey = 1;
}

message UpdateValue {
  bytes userKey = 1;
  bytes value = 2;
}

message Iterator {
  string iteratorId = 1;
}

message Keys {
  string iteratorId = 1;
}

message Values {
  string iteratorId = 1;
}

message RemoveKey {
  bytes userKey = 1;
}

enum HandleState {
  PRE_INIT = 0;
  CREATED = 1;
  INITIALIZED = 2;
  DATA_PROCESSED = 3;
  TIMER_PROCESSED = 4;
  CLOSED = 5;
}

message SetHandleState {
  HandleState state = 1;
}

message TTLConfig {
  int64 durationMs = 1;
}
