agent-infra-sdk 0.1.1

use crate::operation::{OperationHandle, OperationObservation, OperationPoller, OperationProgress};
use crate::transport::{CallOptions, HttpTransport, InfraClientError, ServiceEndpoint};
use agent_workspace_contract::{
    AckResponse, BytesResponse, CloneWorkspaceRequest, CmdOutput, CommandResource, ContentResponse,
    CreateCommandRequest, CreatePreviewRequest, CreateSnapshotRequest, CreateUploadRequest,
    CreateWorkspaceRequest, DirEntry, EntriesResponse, ExecRequest, ExistsResponse,
    KeepAliveRequest, ListDirRequest, MigrateWorkspaceRequest, OperationState, OutputPage,
    PathRequest, PutWorkspaceFileRequest, StoredPreview, UploadChunkRequest, UploadSession,
    WORKSPACE_CAPABILITIES_PATH, WORKSPACE_COMMAND_CANCEL_PATH, WORKSPACE_COMMAND_OUTPUT_PATH,
    WORKSPACE_COMMAND_PATH, WORKSPACE_COMMANDS_PATH, WORKSPACE_CREATE_DIR_ALL_PATH,
    WORKSPACE_EXEC_PATH, WORKSPACE_EXISTS_PATH, WORKSPACE_LEASE_PATH, WORKSPACE_LIST_DIR_PATH,
    WORKSPACE_MIGRATE_PATH, WORKSPACE_OPERATION_CANCEL_PATH, WORKSPACE_OPERATIONS_PATH,
    WORKSPACE_PATH, WORKSPACE_READ_FILE_PATH, WORKSPACE_RECONCILE_PATH, WORKSPACE_REMOVE_FILE_PATH,
    WORKSPACE_RESOURCE_FILE_PATH, WORKSPACE_RESOURCE_SEARCH_PATH, WORKSPACE_RESUME_PATH,
    WORKSPACE_SERVICE_NAME, WORKSPACE_SNAPSHOT_RESTORE_PATH, WORKSPACE_SNAPSHOTS_PATH,
    WORKSPACE_SUSPEND_PATH, WORKSPACE_UPLOAD_PATH, WORKSPACE_UPLOADS_PATH, WORKSPACE_USAGE_PATH,
    WORKSPACE_WRITE_FILE_PATH, WORKSPACES_PATH, WorkspaceActionRequest, WorkspaceCapabilities,
    WorkspaceLease, WorkspaceOperation, WorkspacePage, WorkspaceRecord, WorkspaceResourceUsage,
    WorkspaceSearchRequest, WorkspaceSearchResponse, WriteFileRequest,
};
use async_trait::async_trait;
use reqwest::Client;
use std::sync::Arc;
use std::time::Duration;

/// Client for the Workspace service (`agent-workspace-server`). It owns no
/// workspace state; every wire type comes from `agent-workspace-contract`.
#[derive(Clone, Debug)]
pub struct WorkspaceClient {
    transport: HttpTransport,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct VersionedContent {
    pub content: String,
    pub revision: Option<String>,
}

impl WorkspaceClient {
    pub(crate) fn new_with_endpoint(
        http: Client,
        endpoint: ServiceEndpoint,
        options: crate::ClientOptions,
    ) -> Self {
        let endpoint = endpoint.with_default_credential_audience("workspace");
        Self {
            transport: HttpTransport::new_with_options(
                http,
                WORKSPACE_SERVICE_NAME,
                endpoint,
                options,
            ),
        }
    }

    pub async fn capabilities(&self) -> Result<WorkspaceCapabilities, InfraClientError> {
        self.transport.get_json(WORKSPACE_CAPABILITIES_PATH).await
    }

    pub async fn create_workspace(
        &self,
        request: &CreateWorkspaceRequest,
        idempotency_key: &str,
    ) -> Result<OperationHandle<WorkspaceOperation>, InfraClientError> {
        let operation = self
            .transport
            .post_json_with_options(
                WORKSPACES_PATH,
                request,
                CallOptions::default().idempotency_key(idempotency_key),
            )
            .await?;
        Ok(self.operation_handle(operation))
    }

    pub async fn workspace(&self, workspace_id: &str) -> Result<WorkspaceRecord, InfraClientError> {
        self.workspace_with_options(workspace_id, CallOptions::default())
            .await
    }

    pub async fn workspace_with_options(
        &self,
        workspace_id: &str,
        options: CallOptions,
    ) -> Result<WorkspaceRecord, InfraClientError> {
        self.transport
            .get_json_with_options(&expand(WORKSPACE_PATH, &[workspace_id]), options)
            .await
    }

    pub async fn list_workspaces(
        &self,
        after: Option<&str>,
        limit: Option<usize>,
    ) -> Result<WorkspacePage, InfraClientError> {
        let mut path = format!("{WORKSPACES_PATH}?limit={}", limit.unwrap_or(50).min(200));
        if let Some(after) = after {
            path.push_str("&after=");
            path.push_str(&encode_query_component(after));
        }
        self.transport.get_json(&path).await
    }

    pub async fn suspend_workspace(
        &self,
        workspace_id: &str,
        request: &WorkspaceActionRequest,
        key: &str,
    ) -> Result<OperationHandle<WorkspaceOperation>, InfraClientError> {
        self.workspace_action(WORKSPACE_SUSPEND_PATH, workspace_id, request, key)
            .await
    }

    pub async fn resume_workspace(
        &self,
        workspace_id: &str,
        request: &WorkspaceActionRequest,
        key: &str,
    ) -> Result<OperationHandle<WorkspaceOperation>, InfraClientError> {
        self.workspace_action(WORKSPACE_RESUME_PATH, workspace_id, request, key)
            .await
    }

    pub async fn reconcile_workspace(
        &self,
        workspace_id: &str,
        request: &WorkspaceActionRequest,
        key: &str,
    ) -> Result<OperationHandle<WorkspaceOperation>, InfraClientError> {
        self.workspace_action(WORKSPACE_RECONCILE_PATH, workspace_id, request, key)
            .await
    }

    pub async fn delete_workspace(
        &self,
        workspace_id: &str,
        request: &WorkspaceActionRequest,
        key: &str,
    ) -> Result<OperationHandle<WorkspaceOperation>, InfraClientError> {
        let operation = self
            .transport
            .delete_json_with_options(
                &expand(WORKSPACE_PATH, &[workspace_id]),
                request,
                CallOptions::default().idempotency_key(key),
            )
            .await?;
        Ok(self.operation_handle(operation))
    }

    pub async fn workspace_usage(
        &self,
        workspace_id: &str,
    ) -> Result<WorkspaceResourceUsage, InfraClientError> {
        self.transport
            .get_json(&expand(WORKSPACE_USAGE_PATH, &[workspace_id]))
            .await
    }

    async fn workspace_action(
        &self,
        path: &str,
        workspace_id: &str,
        request: &WorkspaceActionRequest,
        key: &str,
    ) -> Result<OperationHandle<WorkspaceOperation>, InfraClientError> {
        let operation = self
            .transport
            .post_json_with_options(
                &expand(path, &[workspace_id]),
                request,
                CallOptions::default().idempotency_key(key),
            )
            .await?;
        Ok(self.operation_handle(operation))
    }

    pub async fn read_file(&self, path: &str) -> Result<String, InfraClientError> {
        Ok(self.read_file_versioned(path).await?.content)
    }

    /// Read a file from an explicit managed workspace. New consumers should
    /// prefer this over the implicit single-root compatibility API.
    pub async fn read_workspace_file(
        &self,
        workspace_id: &str,
        path: &str,
    ) -> Result<BytesResponse, InfraClientError> {
        self.read_workspace_file_with_options(workspace_id, path, CallOptions::default())
            .await
    }

    pub async fn read_workspace_file_with_options(
        &self,
        workspace_id: &str,
        path: &str,
        options: CallOptions,
    ) -> Result<BytesResponse, InfraClientError> {
        self.transport
            .get_json_with_options(&resource_file_path(workspace_id, path), options)
            .await
    }

    pub async fn put_workspace_file(
        &self,
        workspace_id: &str,
        path: &str,
        request: &PutWorkspaceFileRequest,
        options: CallOptions,
    ) -> Result<Option<String>, InfraClientError> {
        let response: AckResponse = self
            .transport
            .put_json_with_options(&resource_file_path(workspace_id, path), request, options)
            .await?;
        ensure_ack(response)
    }

    pub async fn search_workspace_files(
        &self,
        workspace_id: &str,
        request: &WorkspaceSearchRequest,
    ) -> Result<WorkspaceSearchResponse, InfraClientError> {
        self.search_workspace_files_with_options(workspace_id, request, CallOptions::default())
            .await
    }

    pub async fn search_workspace_files_with_options(
        &self,
        workspace_id: &str,
        request: &WorkspaceSearchRequest,
        options: CallOptions,
    ) -> Result<WorkspaceSearchResponse, InfraClientError> {
        self.transport
            .post_json_with_options(
                &expand(WORKSPACE_RESOURCE_SEARCH_PATH, &[workspace_id]),
                request,
                options.idempotent(true),
            )
            .await
    }

    pub async fn next_workspace_search(
        &self,
        workspace_id: &str,
        request: &WorkspaceSearchRequest,
        current: &WorkspaceSearchResponse,
    ) -> Result<Option<WorkspaceSearchResponse>, InfraClientError> {
        let Some(cursor) = current.next_cursor.as_deref() else {
            return Ok(None);
        };
        let mut next = request.clone();
        next.cursor = Some(cursor.to_string());
        self.search_workspace_files(workspace_id, &next)
            .await
            .map(Some)
    }

    pub async fn read_file_versioned(
        &self,
        path: &str,
    ) -> Result<VersionedContent, InfraClientError> {
        let response: ContentResponse = self
            .transport
            .post_json_idempotent(WORKSPACE_READ_FILE_PATH, &PathRequest { path: path.into() })
            .await?;
        Ok(VersionedContent {
            content: response.content,
            revision: response.revision,
        })
    }

    pub async fn write_file(&self, path: &str, content: &str) -> Result<(), InfraClientError> {
        self.write_file_with_options(path, content, None, CallOptions::default())
            .await
            .map(|_| ())
    }

    pub async fn write_file_if_match(
        &self,
        path: &str,
        content: &str,
        revision: &str,
    ) -> Result<String, InfraClientError> {
        self.write_file_with_options(
            path,
            content,
            Some(revision),
            // A lost successful CAS response cannot be replayed safely without
            // server-side idempotency-key storage: the original revision has
            // already changed. Keep automatic retry disabled by default.
            CallOptions::default(),
        )
        .await?
        .ok_or_else(|| InfraClientError::Protocol {
            service: WORKSPACE_SERVICE_NAME,
            message: "workspace omitted the revision after a conditional write".into(),
        })
    }

    pub async fn write_file_with_options(
        &self,
        path: &str,
        content: &str,
        if_match: Option<&str>,
        options: CallOptions,
    ) -> Result<Option<String>, InfraClientError> {
        let response: AckResponse = self
            .transport
            .post_json_with_options(
                WORKSPACE_WRITE_FILE_PATH,
                &WriteFileRequest {
                    path: path.into(),
                    content: content.into(),
                    if_match: if_match.map(Into::into),
                },
                options,
            )
            .await?;
        ensure_ack(response)
    }

    pub async fn exists(&self, path: &str) -> Result<bool, InfraClientError> {
        let response: ExistsResponse = self
            .transport
            .post_json_idempotent(WORKSPACE_EXISTS_PATH, &PathRequest { path: path.into() })
            .await?;
        Ok(response.exists)
    }

    pub async fn create_dir_all(&self, path: &str) -> Result<(), InfraClientError> {
        let response: AckResponse = self
            .transport
            .post_json(
                WORKSPACE_CREATE_DIR_ALL_PATH,
                &PathRequest { path: path.into() },
            )
            .await?;
        ensure_ack(response).map(|_| ())
    }

    pub async fn list_dir(&self, path: &str) -> Result<Vec<DirEntry>, InfraClientError> {
        let response: EntriesResponse = self
            .transport
            .post_json_idempotent(
                WORKSPACE_LIST_DIR_PATH,
                &ListDirRequest {
                    path: path.into(),
                    include_hidden: false,
                    include_ignored: false,
                },
            )
            .await?;
        Ok(response.entries)
    }

    pub async fn exec(
        &self,
        command: &str,
        cwd: Option<&str>,
    ) -> Result<CmdOutput, InfraClientError> {
        self.transport
            .post_json(
                WORKSPACE_EXEC_PATH,
                &ExecRequest {
                    command: command.into(),
                    cwd: cwd.map(Into::into),
                },
            )
            .await
    }

    pub async fn remove_file(&self, path: &str) -> Result<(), InfraClientError> {
        let response: AckResponse = self
            .transport
            .post_json(
                WORKSPACE_REMOVE_FILE_PATH,
                &PathRequest { path: path.into() },
            )
            .await?;
        ensure_ack(response).map(|_| ())
    }

    pub async fn create_upload(
        &self,
        workspace_id: &str,
        request: &CreateUploadRequest,
        idempotency_key: &str,
    ) -> Result<UploadSession, InfraClientError> {
        self.transport
            .post_json_with_options(
                &expand(WORKSPACE_UPLOADS_PATH, &[workspace_id]),
                request,
                CallOptions::default().idempotency_key(idempotency_key),
            )
            .await
    }

    pub async fn create_preview(
        &self,
        workspace_id: &str,
        request: &CreatePreviewRequest,
        idempotency_key: &str,
    ) -> Result<StoredPreview, InfraClientError> {
        self.transport
            .post_json_with_options(
                &format!(
                    "/internal/v1/workspaces/{}/previews",
                    encode_path_segment(workspace_id)
                ),
                request,
                CallOptions::default().idempotency_key(idempotency_key),
            )
            .await
    }

    pub async fn upload(
        &self,
        workspace_id: &str,
        upload_id: &str,
    ) -> Result<UploadSession, InfraClientError> {
        self.transport
            .get_json(&expand(WORKSPACE_UPLOAD_PATH, &[workspace_id, upload_id]))
            .await
    }

    pub async fn upload_chunk(
        &self,
        workspace_id: &str,
        upload_id: &str,
        request: &UploadChunkRequest,
    ) -> Result<UploadSession, InfraClientError> {
        self.transport
            .put_json_with_options(
                &expand(WORKSPACE_UPLOAD_PATH, &[workspace_id, upload_id]),
                request,
                CallOptions::default(),
            )
            .await
    }

    pub async fn create_command(
        &self,
        workspace_id: &str,
        request: &CreateCommandRequest,
        idempotency_key: &str,
    ) -> Result<CommandResource, InfraClientError> {
        self.create_command_with_options(
            workspace_id,
            request,
            CallOptions::default().idempotency_key(idempotency_key),
        )
        .await
    }

    pub async fn create_command_with_options(
        &self,
        workspace_id: &str,
        request: &CreateCommandRequest,
        options: CallOptions,
    ) -> Result<CommandResource, InfraClientError> {
        self.transport
            .post_json_with_options(
                &expand(WORKSPACE_COMMANDS_PATH, &[workspace_id]),
                request,
                options,
            )
            .await
    }

    pub async fn command(
        &self,
        workspace_id: &str,
        command_id: &str,
    ) -> Result<CommandResource, InfraClientError> {
        self.command_with_options(workspace_id, command_id, CallOptions::default())
            .await
    }

    pub async fn command_with_options(
        &self,
        workspace_id: &str,
        command_id: &str,
        options: CallOptions,
    ) -> Result<CommandResource, InfraClientError> {
        self.transport
            .get_json_with_options(
                &expand(WORKSPACE_COMMAND_PATH, &[workspace_id, command_id]),
                options,
            )
            .await
    }

    pub async fn command_output(
        &self,
        workspace_id: &str,
        command_id: &str,
        after: u64,
        limit: Option<usize>,
    ) -> Result<OutputPage, InfraClientError> {
        self.command_output_with_options(
            workspace_id,
            command_id,
            after,
            limit,
            CallOptions::default(),
        )
        .await
    }

    pub async fn command_output_with_options(
        &self,
        workspace_id: &str,
        command_id: &str,
        after: u64,
        limit: Option<usize>,
        options: CallOptions,
    ) -> Result<OutputPage, InfraClientError> {
        let mut path = expand(WORKSPACE_COMMAND_OUTPUT_PATH, &[workspace_id, command_id]);
        path.push_str(&format!(
            "?after={after}&limit={}",
            limit.unwrap_or(64).min(64)
        ));
        self.transport.get_json_with_options(&path, options).await
    }

    pub async fn next_command_output(
        &self,
        workspace_id: &str,
        command_id: &str,
        current: &OutputPage,
        limit: Option<usize>,
    ) -> Result<Option<OutputPage>, InfraClientError> {
        let Some(cursor) = current.next_cursor else {
            return Ok(None);
        };
        self.command_output(workspace_id, command_id, cursor, limit)
            .await
            .map(Some)
    }

    pub async fn cancel_command(
        &self,
        workspace_id: &str,
        command_id: &str,
    ) -> Result<CommandResource, InfraClientError> {
        self.cancel_command_with_options(
            workspace_id,
            command_id,
            CallOptions::default()
                .idempotency_key(format!("cancel-command:{workspace_id}:{command_id}")),
        )
        .await
    }

    pub async fn cancel_command_with_options(
        &self,
        workspace_id: &str,
        command_id: &str,
        options: CallOptions,
    ) -> Result<CommandResource, InfraClientError> {
        self.transport
            .post_json_with_options(
                &expand(WORKSPACE_COMMAND_CANCEL_PATH, &[workspace_id, command_id]),
                &serde_json::json!({}),
                options,
            )
            .await
    }

    pub async fn create_snapshot(
        &self,
        workspace_id: &str,
        request: &CreateSnapshotRequest,
        idempotency_key: &str,
    ) -> Result<OperationHandle<WorkspaceOperation>, InfraClientError> {
        let operation: WorkspaceOperation = self
            .transport
            .post_json_with_options(
                &expand(WORKSPACE_SNAPSHOTS_PATH, &[workspace_id]),
                request,
                CallOptions::default().idempotency_key(idempotency_key),
            )
            .await?;
        Ok(self.operation_handle(operation))
    }

    pub async fn restore_snapshot(
        &self,
        workspace_id: &str,
        snapshot_id: &str,
        idempotency_key: &str,
    ) -> Result<OperationHandle<WorkspaceOperation>, InfraClientError> {
        let operation: WorkspaceOperation = self
            .transport
            .post_json_with_options(
                &expand(
                    WORKSPACE_SNAPSHOT_RESTORE_PATH,
                    &[workspace_id, snapshot_id],
                ),
                &serde_json::json!({}),
                CallOptions::default().idempotency_key(idempotency_key),
            )
            .await?;
        Ok(self.operation_handle(operation))
    }

    pub async fn clone_workspace(
        &self,
        workspace_id: &str,
        request: &CloneWorkspaceRequest,
        idempotency_key: &str,
    ) -> Result<OperationHandle<WorkspaceOperation>, InfraClientError> {
        let operation: WorkspaceOperation = self
            .transport
            .post_json_with_options(
                &expand(
                    agent_workspace_contract::WORKSPACE_CLONE_PATH,
                    &[workspace_id],
                ),
                request,
                CallOptions::default().idempotency_key(idempotency_key),
            )
            .await?;
        Ok(self.operation_handle(operation))
    }

    pub async fn migrate_workspace(
        &self,
        workspace_id: &str,
        request: &MigrateWorkspaceRequest,
        idempotency_key: &str,
    ) -> Result<OperationHandle<WorkspaceOperation>, InfraClientError> {
        let operation: WorkspaceOperation = self
            .transport
            .post_json_with_options(
                &expand(WORKSPACE_MIGRATE_PATH, &[workspace_id]),
                request,
                CallOptions::default().idempotency_key(idempotency_key),
            )
            .await?;
        Ok(self.operation_handle(operation))
    }

    pub async fn keep_alive(
        &self,
        workspace_id: &str,
        request: &KeepAliveRequest,
    ) -> Result<WorkspaceLease, InfraClientError> {
        self.transport
            .put_json_with_options(
                &expand(WORKSPACE_LEASE_PATH, &[workspace_id]),
                request,
                CallOptions::default(),
            )
            .await
    }

    pub async fn operation(
        &self,
        operation_id: &str,
    ) -> Result<OperationHandle<WorkspaceOperation>, InfraClientError> {
        let operation = WorkspaceOperationPoller {
            transport: self.transport.clone(),
        }
        .poll(operation_id)
        .await?;
        Ok(self.operation_handle(operation))
    }

    fn operation_handle(
        &self,
        operation: WorkspaceOperation,
    ) -> OperationHandle<WorkspaceOperation> {
        OperationHandle::new(
            operation.id.clone(),
            Some(operation),
            Arc::new(WorkspaceOperationPoller {
                transport: self.transport.clone(),
            }),
        )
    }
}

#[derive(Clone, Debug)]
struct WorkspaceOperationPoller {
    transport: HttpTransport,
}

#[async_trait]
impl OperationPoller<WorkspaceOperation> for WorkspaceOperationPoller {
    async fn poll(&self, operation_id: &str) -> Result<WorkspaceOperation, InfraClientError> {
        self.transport
            .get_json(&expand(WORKSPACE_OPERATIONS_PATH, &[operation_id]))
            .await
    }

    async fn cancel(&self, operation_id: &str) -> Result<(), InfraClientError> {
        let operation: WorkspaceOperation = self.poll(operation_id).await?;
        let _: WorkspaceOperation = self
            .transport
            .post_json_with_options(
                &expand(WORKSPACE_OPERATION_CANCEL_PATH, &[operation_id]),
                &WorkspaceActionRequest {
                    expected_version: operation.version,
                    reason: None,
                },
                // Cancellation carries an optimistic version. Until the
                // service stores a cancellation idempotency key, replaying a
                // lost success can manufacture a false version conflict.
                CallOptions::default(),
            )
            .await?;
        Ok(())
    }

    fn observe(&self, operation: &WorkspaceOperation) -> OperationObservation {
        let progress = match operation.state {
            OperationState::Succeeded => OperationProgress::Succeeded,
            OperationState::Failed => OperationProgress::Failed,
            OperationState::Canceled => OperationProgress::Canceled,
            OperationState::Queued | OperationState::Running | OperationState::Canceling => {
                OperationProgress::Pending
            }
        };
        OperationObservation {
            progress,
            next_poll_after: operation.next_poll_after_ms.map(Duration::from_millis),
            error_code: operation.error_code.clone(),
        }
    }
}

fn ensure_ack(response: AckResponse) -> Result<Option<String>, InfraClientError> {
    if response.ok {
        Ok(response.revision)
    } else {
        Err(InfraClientError::Protocol {
            service: WORKSPACE_SERVICE_NAME,
            message: "operation returned ok=false".into(),
        })
    }
}

fn expand(template: &str, values: &[&str]) -> String {
    let mut result = template.to_string();
    for value in values {
        let Some(start) = result.find('{') else {
            break;
        };
        let Some(relative_end) = result[start..].find('}') else {
            break;
        };
        result.replace_range(start..=start + relative_end, &encode_path_segment(value));
    }
    result
}

fn resource_file_path(workspace_id: &str, path: &str) -> String {
    let encoded_path = path
        .trim_start_matches('/')
        .split('/')
        .map(encode_path_segment)
        .collect::<Vec<_>>()
        .join("/");
    expand(WORKSPACE_RESOURCE_FILE_PATH, &[workspace_id]).replace("{*filePath}", &encoded_path)
}

fn encode_path_segment(value: &str) -> String {
    let mut encoded = String::with_capacity(value.len());
    for byte in value.bytes() {
        if byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'.' | b'_' | b'~') {
            encoded.push(char::from(byte));
        } else {
            use std::fmt::Write as _;
            write!(&mut encoded, "%{byte:02X}").expect("writing to String cannot fail");
        }
    }
    encoded
}

fn encode_query_component(value: &str) -> String {
    encode_path_segment(value)
}