// Copyright 2026 Google LLC
//
// Licensed 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 ax.v1alpha1;

import "google/protobuf/struct.proto";
import "google/protobuf/timestamp.proto";

option go_package = "github.com/google/ax/pkg/apis/v1alpha1";

// AX defines the core control plane gRPC service for autonomous agent orchestration.
service AX {
  // Manifests are parsed by the client (see `ax apply`) and submitted through the
  // typed Update* RPCs below; the server never receives raw YAML.

  // Tasks
  rpc GetTask(GetTaskRequest) returns (Task);
  rpc ListTasks(ListTasksRequest) returns (ListTasksResponse);
  rpc CreateTask(CreateTaskRequest) returns (Task);
  rpc DeleteTask(DeleteTaskRequest) returns (DeleteTaskResponse);
  rpc SuspendTask(SuspendTaskRequest) returns (Task);
  rpc ResumeTask(ResumeTaskRequest) returns (Task);
  rpc WatchTask(WatchTaskRequest) returns (stream WatchTaskResponse);


  // Workspaces
  rpc GetWorkspace(GetWorkspaceRequest) returns (Workspace);
  rpc ListWorkspaces(ListWorkspacesRequest) returns (ListWorkspacesResponse);
  rpc UpdateWorkspace(UpdateWorkspaceRequest) returns (Workspace);
  rpc DeleteWorkspace(DeleteWorkspaceRequest) returns (DeleteWorkspaceResponse);

  // Models
  rpc GetModel(GetModelRequest) returns (Model);
  rpc ListModels(ListModelsRequest) returns (ListModelsResponse);
  rpc UpdateModel(UpdateModelRequest) returns (Model);
  rpc DeleteModel(DeleteModelRequest) returns (DeleteModelResponse);

}

// ObjectMeta is metadata for AX resources.
message ObjectMeta {
  string name = 1;
  string atespace = 2;
  google.protobuf.Timestamp creation_timestamp = 3;
}

// --- Task ---

message Task {
  string api_version = 1;
  string kind = 2;
  ObjectMeta metadata = 3;
  TaskSpec spec = 4;
  TaskStatus status = 5;
}

message TaskSpec {
  string image = 1;
  repeated string command = 2;
  repeated EnvVar env = 3;
  ResourceReqs resources = 4;
  // workspaces binds one or more Workspaces, each mounted at its own path
  // under /workspace. The first entry is the task command's working directory.
  repeated WorkspaceRef workspaces = 5;
  // debug enables the in-container guest services (process execution and file
  // access) that back `ax ssh`. Off by default.
  bool debug = 6;
}

message EnvVar {
  string name = 1;
  string value = 2;
}

message ResourceReqs {
  ResourceList requests = 1;
  ResourceList limits = 2;
}

message ResourceList {
  string cpu = 1;
  string memory = 2;
}

message WorkspaceRef {
  string name = 1;
  string path = 2;
  string goal = 3;
}

message TaskStatus {
  string phase = 1;
  string id = 2;
  string actor = 3;
  // json_name keeps the established `workerIP` spelling in YAML and JSON.
  string worker_ip = 4 [json_name = "workerIP"];
  PendingApproval pending_approval = 5;
  UsageStats usage = 6;
  repeated Condition conditions = 7;
}

message PendingApproval {
  string id = 1;
  string action = 2;
  google.protobuf.Timestamp requested_at = 3;
}

message UsageStats {
  int32 prompt_tokens = 1;
  int32 completion_tokens = 2;
}

message Condition {
  string type = 1;
  string status = 2;
  google.protobuf.Timestamp last_transition_time = 3;
  string reason = 4;
  string message = 5;
}

// --- Workspace ---

message Workspace {
  string api_version = 1;
  string kind = 2;
  ObjectMeta metadata = 3;
  WorkspaceSpec spec = 4;
}

message WorkspaceSpec {
  repeated GitRepo git = 1;
  repeated File files = 2;
  MCPConfig mcp = 3;
  SkillsConfig skills = 4;
}

message GitRepo {
  string name = 1;
  string repo = 2;
  string branch = 3;
  string dir = 4;
  int32 depth = 5;
}

message File {
  string path = 1;
  string content = 2;
}

message MCPConfig {
  repeated MCPRegistry registries = 1;
  repeated MCPServer servers = 2;
}

message MCPRegistry {
  string provider = 1;
  string project = 2;
  string query = 3;
  repeated MCPServer servers = 4;
}

message MCPServer {
  string name = 1;
  string endpoint = 2;
  string command = 3;
  repeated string args = 4;
}

message SkillsConfig {
  repeated SkillRegistry registries = 1;
  string path = 2;
}

message SkillRegistry {
  string provider = 1;
  string project = 2;
  string query = 3;
}

// --- Model ---

message Model {
  string api_version = 1;
  string kind = 2;
  ObjectMeta metadata = 3;
  ModelSpec spec = 4;
}

message ModelSpec {
  string provider = 1;
  string model = 2;
  // Fields 3-5 were typed temperature, max_tokens, and system_instruction;
  // provider settings now live in the free-form `parameters` map.
  reserved 3, 4, 5;
  reserved "temperature", "max_tokens", "system_instruction";
  SecretKeyRef secret_key = 6;
  // parameters are provider-specific generation settings passed through to the
  // model API as-is, for example temperature or maxOutputTokens for Gemini.
  google.protobuf.Struct parameters = 7;
}

message SecretKeyRef {
  string name = 1;
  string key = 2;
}

// --- RPC Request & Response Messages ---

// Tasks
message GetTaskRequest {
  string atespace = 1;
  string name = 2;
}

message ListTasksRequest {
  string atespace = 1;
  int64 limit = 2;
  int64 offset = 3;
}

message ListTasksResponse {
  repeated Task tasks = 1;
}

message CreateTaskRequest {
  Task task = 1;
}

message DeleteTaskRequest {
  string atespace = 1;
  string name = 2;
}

message DeleteTaskResponse {}

message SuspendTaskRequest {
  string atespace = 1;
  string name = 2;
}

message ResumeTaskRequest {
  string atespace = 1;
  string name = 2;
}

message WatchTaskRequest {
  string atespace = 1;
  string name = 2;
}

message WatchTaskResponse {
  Task task = 1;
  string action = 2;
}

// Workspaces
message GetWorkspaceRequest {
  string atespace = 1;
  string name = 2;
}

message ListWorkspacesRequest {
  string atespace = 1;
}

message ListWorkspacesResponse {
  repeated Workspace workspaces = 1;
}

message UpdateWorkspaceRequest {
  Workspace workspace = 1;
}

message DeleteWorkspaceRequest {
  string atespace = 1;
  string name = 2;
}

message DeleteWorkspaceResponse {}

// Models
message GetModelRequest {
  string atespace = 1;
  string name = 2;
}

message ListModelsRequest {
  string atespace = 1;
}

message ListModelsResponse {
  repeated Model models = 1;
}

message UpdateModelRequest {
  Model model = 1;
}

message DeleteModelRequest {
  string atespace = 1;
  string name = 2;
}

message DeleteModelResponse {}