syntax = "proto3";

package skytorch.torch;

// Tensor management service for SkyTorch PyTorch backend
service Service {
  // Tensor lifecycle management
  rpc UpdateTensor(stream TensorChunk) returns (TensorResponse);
  rpc GetTensor(GetTensorRequest) returns (stream TensorChunk);
  rpc CopyTensor(CopyTensorRequest) returns (TensorResponse);
  rpc DeleteTensors(DeleteTensorsRequest) returns (TensorResponse);

  // ATen operation execution
  rpc ExecuteAtenOperation(ExecuteAtenRequest) returns (ExecuteAtenResponse);

  // Bidirectional streaming for low-latency operations
  rpc StreamOperations(stream StreamRequest) returns (stream StreamResponse);

  // Execute a function on the server
  rpc ExecuteFunction(ExecuteFunctionRequest) returns (stream ExecuteFunctionEvent);
}

// Request to delete tensors by ID
message DeleteTensorsRequest {
  repeated uint64 tensor_ids = 1;
}

// Tensor metadata for auto-creation on the server
message TensorMetadata {
  uint64 tensor_id = 1;
  repeated int64 shape = 2;
  string dtype = 3;
  int64 nbytes = 4;
  string device_type = 5;
  repeated int64 stride = 6;
  int64 storage_offset = 7;
  int32 device_index = 8;
  optional uint64 tensor_ref = 9;  // Reference to base tensor for views
}

// Represents a chunk of serialized tensor data
message TensorChunk {
  // Tensor ID
  uint64 tensor_id = 1;

  // Chunk sequence number (for ordering)
  uint32 chunk_number = 2;

  // Binary data chunk
  bytes data = 3;

  // Total number of chunks for this tensor (set in first chunk)
  uint32 total_chunks = 4;

  // Tensor info (set on first chunk only)
  repeated int64 shape = 5;
  repeated int64 stride = 6;
  int64 storage_offset = 7;
  string dtype = 8;

  // Total serialized size in bytes (set in first chunk, for pre-allocation)
  uint64 total_bytes = 9;

  // Tensor metadata for auto-creation on first chunk
  TensorMetadata metadata = 10;
}

// Response after receiving tensors
message TensorResponse {
  bool success = 1;
  string message = 2;
}

// Request to get a scalar value from a tensor
message GetScalarRequest {
  uint64 tensor_id = 1;
  repeated TensorMetadata tensor_metadata = 2;
}

// Response containing a scalar value
message GetScalarResponse {
  oneof value {
    double float_value = 1;
    int64 int_value = 2;
    bool bool_value = 3;
  }
}

// Request to create a tensor on the server
message CreateTensorRequest {
  uint64 tensor_id = 1;
  repeated int64 shape = 2;
  string dtype = 3;
  int64 nbytes = 4;
  string device_type = 5;
  repeated int64 stride = 6;
  int64 storage_offset = 7;
  int32 device_index = 8;
  optional uint64 tensor_ref = 9;  // Reference to base tensor for views
}

// Request to get tensor data from storage
message GetTensorRequest {
  uint64 tensor_id = 1;
  repeated int64 shape = 2;
  string dtype = 3;
  repeated int64 stride = 4;
  int64 storage_offset = 5;
  TensorMetadata metadata = 6;  // For auto-creation of tensor on server
}

// Request to copy tensor data on the server
message CopyTensorRequest {
  uint64 src_tensor_id = 1;
  uint64 dst_tensor_id = 2;
  int64 src_offset = 3;
  int64 dst_offset = 4;
  int64 num_bytes = 5;
  TensorMetadata src_metadata = 6;  // For auto-creation of source tensor
  TensorMetadata dst_metadata = 7;  // For auto-creation of destination tensor
}

// Reference to a tensor on the server
message TensorReference {
  uint64 tensor_id = 1;
}

// Argument for ATen operations (can be tensor, scalar, or nested structure)
message AtenArgument {
  oneof value {
    TensorReference tensor = 1;
    double scalar_float = 2;
    int64 scalar_int = 3;
    bool scalar_bool = 4;
    string scalar_string = 5;
    AtenArgumentList list_value = 6;
    bool none_value = 7;  // Represents Python None
    string scalar_dtype = 8;  // torch.dtype as string (e.g., "torch.float32")
    string scalar_memory_format = 9;  // torch.memory_format as string (e.g., "torch.contiguous_format")
    string scalar_layout = 10;  // torch.layout as string (e.g., "torch.strided")
  }
}

// List of AtenArguments (for nested lists/tuples)
message AtenArgumentList {
  repeated AtenArgument values = 1;
  bool is_tuple = 2;  // Preserve tuple vs list distinction
}

// Request to execute an ATen operation
message ExecuteAtenRequest {
  string op_name = 1;
  repeated AtenArgument args = 2;  // Positional arguments (preserves order)
  repeated TensorReference outputs = 3;
  map<string, AtenArgument> kwargs = 4;
  repeated TensorMetadata tensor_metadata = 5;   // Input tensors to auto-create
  repeated TensorMetadata output_metadata = 6;   // Output tensors to auto-create
}

// Response after executing an ATen operation
message ExecuteAtenResponse {
  bool success = 1;
  string message = 2;
  repeated TensorReference output_tensors = 3;
}

// Request to update tensor data (upload from client)
message UpdateTensorRequest {
  uint64 tensor_id = 1;
  bytes data = 2;  // Serialized tensor data
  repeated int64 shape = 3;
  string dtype = 4;
  repeated int64 stride = 5;
  int64 storage_offset = 6;
  TensorMetadata metadata = 7;  // For auto-creation
}

// Response for GetTensor containing tensor data
message GetTensorResponse {
  bool success = 1;
  string message = 2;
  bytes data = 3;  // Serialized tensor data
  repeated int64 shape = 4;
  string dtype = 5;
  repeated int64 stride = 6;
  int64 storage_offset = 7;
}

// Request to execute a function on the server
message ExecuteFunctionRequest {
  bytes args = 1;               // pickle'd args tuple
  bytes kwargs = 2;             // pickle'd kwargs dict
  string callable_source = 3;   // function source code from inspect.getsource()
  string callable_name = 4;     // function name (fn.__name__)
  bool retain_model = 5;        // Keep model alive on server after tensor extraction
}

// Information about a tensor created on the server
message RemoteTensorInfo {
  string name = 1;
  int64 storage_id = 2;
  repeated int64 shape = 3;
  string dtype = 4;
  repeated int64 stride = 5;
  int64 storage_offset = 6;
  int64 storage_nbytes = 7;
  string device_type = 8;
  int32 device_index = 9;
}

// Response from executing a function on the server
message ExecuteFunctionResponse {
  bool success = 1;
  string error_message = 2;
  repeated RemoteTensorInfo tensors = 3;
  uint64 model_id = 4;          // Server-assigned model ID (when retain_model=true)
}

message ExecuteFunctionLog {
  string stream = 1;  // "stdout" or "stderr"
  string text = 2;
}

message ExecuteFunctionEvent {
  oneof event {
    ExecuteFunctionLog log = 1;
    ExecuteFunctionResponse result = 2;
  }
}

// Request to register tensor ID mappings for server-created tensors
message RegisterTensorsRequest {
  repeated TensorRegistration registrations = 1;
}

// Maps a server-assigned storage_id to the client-computed tensor_id
message TensorRegistration {
  int64 storage_id = 1;
  uint64 tensor_id = 2;
}

// Request to execute a module's forward method on a retained model
message ExecuteModuleForwardRequest {
  uint64 model_id = 1;                        // Server-assigned model ID
  string module_path = 2;                      // Dotted path to module (e.g., "layers.0.experts")
  repeated uint64 input_tensor_ids = 3;        // Tensor IDs for positional inputs
  repeated TensorMetadata input_metadata = 4;  // For auto-creation of input tensors
  repeated uint64 output_tensor_ids = 5;       // Pre-allocated output tensor IDs (fire-and-forget mode)
  repeated TensorMetadata output_metadata = 6; // Output tensors to auto-create (fire-and-forget mode)
}

// Response from executing a module's forward method
message ExecuteModuleForwardResponse {
  bool success = 1;
  string error_message = 2;
  repeated RemoteTensorInfo output_tensors = 3;
}

// Request to release a retained model
message ReleaseModelRequest {
  uint64 model_id = 1;
}

// Batch of execute_aten requests for reduced per-operation overhead
message BatchedExecuteAtenRequest {
  repeated ExecuteAtenRequest operations = 1;
}

// Request wrapper for bidirectional streaming
// All operations go through this single channel to ensure ordering
message StreamRequest {
  oneof request {
    ExecuteAtenRequest execute_aten = 1;
    DeleteTensorsRequest delete_tensors = 2;
    CopyTensorRequest copy_tensor = 3;
    UpdateTensorRequest update_tensor = 4;
    GetTensorRequest get_tensor = 5;
    RegisterTensorsRequest register_tensors = 6;
    BatchedExecuteAtenRequest batched_execute_aten = 7;
    bytes raw_execute_aten = 8;           // Binary-serialized execute_aten (C++ fast path)
    bytes raw_batched_execute_aten = 9;   // Multiple binary-serialized ops concatenated
    GetScalarRequest get_scalar = 10;    // Get scalar value from a tensor
    ExecuteModuleForwardRequest execute_module_forward = 14;  // Execute module forward on retained model
    ReleaseModelRequest release_model = 15;                   // Release a retained model
  }

  // Chunking support for large payloads
  uint32 chunk_number = 11;      // 0-indexed chunk sequence
  uint32 total_chunks = 12;      // Total chunks (0 or 1 = single message)
  uint64 total_bytes = 13;       // Total data size (set on first chunk)
}

// Response wrapper for bidirectional streaming
message StreamResponse {
  bool success = 1;
  string error_message = 2;

  oneof response {
    ExecuteAtenResponse execute_aten = 3;
    TensorResponse delete_tensors = 4;
    TensorResponse copy_tensor = 5;
    TensorResponse update_tensor = 6;
    GetTensorResponse get_tensor = 7;
    TensorResponse register_tensors = 8;
    GetScalarResponse get_scalar = 9;
    ExecuteModuleForwardResponse execute_module_forward = 14;
  }

  // Chunking support for large responses (get_tensor)
  uint32 chunk_number = 10;
  uint32 total_chunks = 11;

  // Server-side timing for sync_wait decomposition (profiling)
  uint64 server_backlog_ns = 12;
  uint64 server_handle_ns = 13;
}
