agent-infra-sdk 0.1.1

use crate::transport::{
    CallOptions, ClientOptions, HttpTransport, InfraClientError, ServiceEndpoint,
};
use agent_context_contract::{
    AppendMessagesRequest, AppendResponse, CONTEXT_ADVANCED_OPERATIONS_PATH, CONTEXT_CORRECT_PATH,
    CONTEXT_EXPIRE_PATH, CONTEXT_EXPORT_ARTIFACT_CHUNK_PATH, CONTEXT_EXPORT_ARTIFACT_PATH,
    CONTEXT_EXPORT_PATH, CONTEXT_IMPORT_PATH, CONTEXT_INSIGHT_GET_PATH, CONTEXT_INSIGHTS_PATH,
    CONTEXT_MERGE_PATH, CONTEXT_OPERATIONS_PATH, CONTEXT_RECALL_PATH, CONVERSATION_DELETE_PATH,
    ContextDerivation, ContextExportArtifact, ContextExportArtifactChunk, ContextInsight,
    ContextMaintenanceStatus, ContextOperation, ContextOperationState, ContextSnapshot,
    ConversationRecord, ConversationScope, CorrectContextInsightRequest,
    CreateContextDerivationRequest, CreateContextSnapshotRequest, DERIVATIONS_PATH,
    DeleteConversationRequest, DeletionOperation, ExpireContextInsightRequest,
    ExportContextRequest, GetContextExportArtifactChunkRequest, GetContextExportArtifactRequest,
    GetContextInsightRequest, GetContextOperationRequest, GetContextSnapshotRequest,
    GetDeletionOperationRequest, GetModelInvocationRequest, ImportContextRequest,
    MAINTENANCE_STATUS_PATH, MESSAGES_APPEND_PATH, MESSAGES_PAGE_PATH, MESSAGES_RECENT_PATH,
    MESSAGES_SEARCH_PATH, MODEL_INVOCATION_GET_PATH, MODEL_INVOCATIONS_PATH,
    MergeContextInsightsRequest, MessagePage, MessagesResponse, ModelInvocation,
    PageMessagesRequest, PutContextInsightRequest, PutModelInvocationRequest, RETENTION_PATH,
    RebuildVectorProjectionRequest, RecallContextRequest, RecallContextResponse,
    RecentMessagesRequest, RetentionRequest, SNAPSHOT_GET_PATH, SNAPSHOTS_PATH,
    SearchMessagesRequest, VECTOR_REBUILD_PATH,
};
use agent_knowledge_contract::{
    KnowledgeCapabilities, KnowledgeQueryRequest, KnowledgeQueryResponse,
    ResolveKnowledgeResourcesRequest, ResolveKnowledgeResourcesResponse,
};
use async_trait::async_trait;
use reqwest::Client;
use std::sync::Arc;
use std::time::Duration;

use crate::{OperationHandle, OperationObservation, OperationPoller, OperationProgress};

/// Client for the Context service (`agent-context-server`).
#[derive(Clone, Debug)]
pub struct ContextClient {
    transport: HttpTransport,
}

impl ContextClient {
    pub(crate) fn new_with_endpoint(
        http: Client,
        endpoint: ServiceEndpoint,
        options: ClientOptions,
    ) -> Self {
        let endpoint = endpoint.with_default_credential_audience(agent_context_contract::AUDIENCE);
        Self {
            transport: HttpTransport::new_with_options(
                http,
                agent_context_contract::SERVICE_NAME,
                endpoint,
                options,
            ),
        }
    }

    /// Append messages to a conversation. Idempotent: duplicate IDs are skipped.
    pub async fn append_messages(
        &self,
        messages: Vec<ConversationRecord>,
    ) -> Result<AppendResponse, InfraClientError> {
        self.append_messages_with_options(messages, CallOptions::default())
            .await
    }

    pub async fn append_messages_with_options(
        &self,
        messages: Vec<ConversationRecord>,
        options: CallOptions,
    ) -> Result<AppendResponse, InfraClientError> {
        self.transport
            .post_json_with_options(
                MESSAGES_APPEND_PATH,
                &AppendMessagesRequest { messages },
                options.idempotent(true),
            )
            .await
    }

    pub async fn create_snapshot(
        &self,
        request: &CreateContextSnapshotRequest,
    ) -> Result<ContextSnapshot, InfraClientError> {
        self.create_snapshot_with_options(request, CallOptions::default())
            .await
    }

    pub async fn create_snapshot_with_options(
        &self,
        request: &CreateContextSnapshotRequest,
        options: CallOptions,
    ) -> Result<ContextSnapshot, InfraClientError> {
        self.transport
            .post_json_with_options(SNAPSHOTS_PATH, request, options.idempotent(true))
            .await
    }

    pub async fn get_snapshot(
        &self,
        snapshot_id: &str,
    ) -> Result<ContextSnapshot, InfraClientError> {
        self.transport
            .post_json_idempotent(
                SNAPSHOT_GET_PATH,
                &GetContextSnapshotRequest {
                    snapshot_id: snapshot_id.to_string(),
                },
            )
            .await
    }

    pub async fn create_derivation(
        &self,
        request: &CreateContextDerivationRequest,
    ) -> Result<ContextDerivation, InfraClientError> {
        self.create_derivation_with_options(request, CallOptions::default())
            .await
    }

    pub async fn create_derivation_with_options(
        &self,
        request: &CreateContextDerivationRequest,
        options: CallOptions,
    ) -> Result<ContextDerivation, InfraClientError> {
        self.transport
            .post_json_with_options(DERIVATIONS_PATH, request, options.idempotent(true))
            .await
    }

    pub async fn put_model_invocation(
        &self,
        request: &PutModelInvocationRequest,
    ) -> Result<ModelInvocation, InfraClientError> {
        self.put_model_invocation_with_options(request, CallOptions::default())
            .await
    }

    pub async fn put_model_invocation_with_options(
        &self,
        request: &PutModelInvocationRequest,
        options: CallOptions,
    ) -> Result<ModelInvocation, InfraClientError> {
        self.transport
            .post_json_with_options(MODEL_INVOCATIONS_PATH, request, options.idempotent(true))
            .await
    }

    pub async fn get_model_invocation(
        &self,
        invocation_id: &str,
    ) -> Result<ModelInvocation, InfraClientError> {
        self.transport
            .post_json_idempotent(
                MODEL_INVOCATION_GET_PATH,
                &GetModelInvocationRequest {
                    invocation_id: invocation_id.to_string(),
                },
            )
            .await
    }

    pub async fn knowledge_capabilities(&self) -> Result<KnowledgeCapabilities, InfraClientError> {
        self.transport
            .get_json(agent_knowledge_contract::KNOWLEDGE_CAPABILITIES_PATH)
            .await
    }

    pub async fn query_knowledge(
        &self,
        request: &KnowledgeQueryRequest,
    ) -> Result<KnowledgeQueryResponse, InfraClientError> {
        self.transport
            .post_json_idempotent(agent_knowledge_contract::KNOWLEDGE_QUERY_PATH, request)
            .await
    }

    pub async fn resolve_knowledge_resources(
        &self,
        request: &ResolveKnowledgeResourcesRequest,
    ) -> Result<ResolveKnowledgeResourcesResponse, InfraClientError> {
        self.transport
            .post_json_idempotent(agent_knowledge_contract::KNOWLEDGE_RESOLVE_PATH, request)
            .await
    }

    /// Get the most recent messages in a conversation.
    pub async fn recent_messages(
        &self,
        conversation_id: &str,
        limit: Option<usize>,
        roles: &[&str],
    ) -> Result<Vec<ConversationRecord>, InfraClientError> {
        self.recent_messages_with_options(conversation_id, limit, roles, CallOptions::default())
            .await
    }

    pub async fn recent_messages_with_options(
        &self,
        conversation_id: &str,
        limit: Option<usize>,
        roles: &[&str],
        options: CallOptions,
    ) -> Result<Vec<ConversationRecord>, InfraClientError> {
        let request = RecentMessagesRequest {
            scope: ConversationScope::conversation(conversation_id),
            limit,
            roles: roles.iter().map(|role| (*role).to_string()).collect(),
        };
        let response: MessagesResponse = self
            .transport
            .post_json_with_options(MESSAGES_RECENT_PATH, &request, options.idempotent(true))
            .await?;
        Ok(response.items)
    }

    /// Search messages by text content.
    pub async fn search_messages(
        &self,
        conversation_id: &str,
        query: &str,
        role: Option<&str>,
        limit: Option<usize>,
    ) -> Result<Vec<ConversationRecord>, InfraClientError> {
        let request = SearchMessagesRequest {
            scope: ConversationScope::conversation(conversation_id),
            query: Some(query.to_string()),
            role: role.map(ToOwned::to_owned),
            limit,
        };
        let response: MessagesResponse = self
            .transport
            .post_json_idempotent(MESSAGES_SEARCH_PATH, &request)
            .await?;
        Ok(response.items)
    }

    /// Fetch one bounded cursor page. The SDK never implicitly drains all pages.
    pub async fn page_messages(
        &self,
        request: &PageMessagesRequest,
    ) -> Result<MessagePage, InfraClientError> {
        self.transport
            .post_json_idempotent(MESSAGES_PAGE_PATH, request)
            .await
    }

    /// Fetch the next bounded page while preserving every original filter.
    pub async fn next_message_page(
        &self,
        request: &PageMessagesRequest,
        current: &MessagePage,
    ) -> Result<Option<MessagePage>, InfraClientError> {
        let Some(cursor) = current.next_cursor.clone() else {
            return Ok(None);
        };
        let mut next = request.clone();
        next.cursor = Some(cursor);
        self.page_messages(&next).await.map(Some)
    }

    pub async fn delete_conversation(
        &self,
        request: &DeleteConversationRequest,
    ) -> Result<DeletionOperation, InfraClientError> {
        self.transport
            .post_json(CONVERSATION_DELETE_PATH, request)
            .await
    }

    pub async fn delete_conversation_operation(
        &self,
        request: &DeleteConversationRequest,
    ) -> Result<OperationHandle<DeletionOperation>, InfraClientError> {
        let initial = self.delete_conversation(request).await?;
        Ok(OperationHandle::new(
            initial.id.clone(),
            Some(initial),
            Arc::new(ContextDeletionPoller {
                transport: self.transport.clone(),
            }),
        ))
    }

    pub async fn get_deletion_operation(
        &self,
        operation_id: &str,
    ) -> Result<DeletionOperation, InfraClientError> {
        self.transport
            .post_json_idempotent(
                CONTEXT_OPERATIONS_PATH,
                &GetDeletionOperationRequest {
                    operation_id: operation_id.to_string(),
                },
            )
            .await
    }

    pub async fn apply_retention(
        &self,
        request: &RetentionRequest,
    ) -> Result<DeletionOperation, InfraClientError> {
        self.transport.post_json(RETENTION_PATH, request).await
    }

    pub async fn maintenance_status(&self) -> Result<ContextMaintenanceStatus, InfraClientError> {
        self.transport.get_json(MAINTENANCE_STATUS_PATH).await
    }

    pub async fn put_context_insight(
        &self,
        request: &PutContextInsightRequest,
        idempotency_key: &str,
    ) -> Result<ContextInsight, InfraClientError> {
        self.transport
            .post_json_with_options(
                CONTEXT_INSIGHTS_PATH,
                request,
                CallOptions::default().idempotency_key(idempotency_key),
            )
            .await
    }

    pub async fn get_context_insight(
        &self,
        entry_id: &str,
    ) -> Result<ContextInsight, InfraClientError> {
        self.transport
            .post_json_idempotent(
                CONTEXT_INSIGHT_GET_PATH,
                &GetContextInsightRequest {
                    entry_id: entry_id.to_string(),
                },
            )
            .await
    }

    pub async fn recall(
        &self,
        request: &RecallContextRequest,
    ) -> Result<RecallContextResponse, InfraClientError> {
        self.transport
            .post_json_idempotent(CONTEXT_RECALL_PATH, request)
            .await
    }

    pub async fn merge_context_insights(
        &self,
        request: &MergeContextInsightsRequest,
        idempotency_key: &str,
    ) -> Result<ContextInsight, InfraClientError> {
        self.transport
            .post_json_with_options(
                CONTEXT_MERGE_PATH,
                request,
                CallOptions::default().idempotency_key(idempotency_key),
            )
            .await
    }

    pub async fn correct_context_insight(
        &self,
        request: &CorrectContextInsightRequest,
    ) -> Result<ContextInsight, InfraClientError> {
        self.transport
            .post_json_with_options(
                CONTEXT_CORRECT_PATH,
                request,
                CallOptions::default().idempotency_key(format!(
                    "correct:{}:{}",
                    request.entry_id, request.expected_version
                )),
            )
            .await
    }

    pub async fn expire_context_insight(
        &self,
        request: &ExpireContextInsightRequest,
    ) -> Result<ContextInsight, InfraClientError> {
        self.transport
            .post_json_with_options(
                CONTEXT_EXPIRE_PATH,
                request,
                CallOptions::default().idempotency_key(format!(
                    "expire:{}:{}",
                    request.entry_id, request.expected_version
                )),
            )
            .await
    }

    pub async fn export_context(
        &self,
        request: &ExportContextRequest,
        idempotency_key: &str,
    ) -> Result<OperationHandle<ContextOperation>, InfraClientError> {
        self.start_context_operation(CONTEXT_EXPORT_PATH, request, idempotency_key)
            .await
    }

    pub async fn import_context(
        &self,
        request: &ImportContextRequest,
    ) -> Result<OperationHandle<ContextOperation>, InfraClientError> {
        let key = request.idempotency_key.clone();
        self.start_context_operation(CONTEXT_IMPORT_PATH, request, &key)
            .await
    }

    pub async fn get_export_artifact(
        &self,
        artifact_id: &str,
    ) -> Result<ContextExportArtifact, InfraClientError> {
        self.transport
            .post_json_idempotent(
                CONTEXT_EXPORT_ARTIFACT_PATH,
                &GetContextExportArtifactRequest {
                    artifact_id: artifact_id.to_string(),
                },
            )
            .await
    }

    pub async fn export_chunk(
        &self,
        artifact_id: &str,
        offset: u64,
        limit: usize,
    ) -> Result<ContextExportArtifactChunk, InfraClientError> {
        self.transport
            .post_json_idempotent(
                CONTEXT_EXPORT_ARTIFACT_CHUNK_PATH,
                &GetContextExportArtifactChunkRequest {
                    artifact_id: artifact_id.to_string(),
                    offset,
                    limit: limit.clamp(1, 1024 * 1024),
                },
            )
            .await
    }

    pub async fn next_export_chunk(
        &self,
        current: &ContextExportArtifactChunk,
        limit: usize,
    ) -> Result<Option<ContextExportArtifactChunk>, InfraClientError> {
        if current.eof {
            return Ok(None);
        }
        self.export_chunk(&current.artifact_id, current.next_offset, limit)
            .await
            .map(Some)
    }

    pub async fn rebuild_vector_projection(
        &self,
        request: &RebuildVectorProjectionRequest,
        idempotency_key: &str,
    ) -> Result<OperationHandle<ContextOperation>, InfraClientError> {
        self.start_context_operation(VECTOR_REBUILD_PATH, request, idempotency_key)
            .await
    }

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

    async fn start_context_operation<Request>(
        &self,
        path: &str,
        request: &Request,
        idempotency_key: &str,
    ) -> Result<OperationHandle<ContextOperation>, InfraClientError>
    where
        Request: serde::Serialize + ?Sized,
    {
        let operation: ContextOperation = self
            .transport
            .post_json_with_options(
                path,
                request,
                CallOptions::default().idempotency_key(idempotency_key),
            )
            .await?;
        Ok(self.context_operation_handle(operation))
    }

    fn context_operation_handle(
        &self,
        operation: ContextOperation,
    ) -> OperationHandle<ContextOperation> {
        OperationHandle::new(
            operation.id.clone(),
            Some(operation),
            Arc::new(ContextOperationPoller {
                transport: self.transport.clone(),
            }),
        )
    }
}

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

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

#[async_trait]
impl OperationPoller<ContextOperation> for ContextOperationPoller {
    async fn poll(&self, operation_id: &str) -> Result<ContextOperation, InfraClientError> {
        self.transport
            .post_json_idempotent(
                CONTEXT_ADVANCED_OPERATIONS_PATH,
                &GetContextOperationRequest {
                    operation_id: operation_id.to_string(),
                },
            )
            .await
    }

    async fn cancel(&self, _operation_id: &str) -> Result<(), InfraClientError> {
        Err(InfraClientError::Protocol {
            service: agent_context_contract::SERVICE_NAME,
            message: "this context operation does not support remote cancellation".into(),
        })
    }

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

#[async_trait]
impl OperationPoller<DeletionOperation> for ContextDeletionPoller {
    async fn poll(&self, operation_id: &str) -> Result<DeletionOperation, InfraClientError> {
        self.transport
            .post_json_idempotent(
                CONTEXT_OPERATIONS_PATH,
                &GetDeletionOperationRequest {
                    operation_id: operation_id.to_string(),
                },
            )
            .await
    }

    async fn cancel(&self, _operation_id: &str) -> Result<(), InfraClientError> {
        Err(InfraClientError::Protocol {
            service: agent_context_contract::SERVICE_NAME,
            message: "accepted deletion operations are irreversible".into(),
        })
    }

    fn observe(&self, operation: &DeletionOperation) -> OperationObservation {
        let progress = match operation.state.as_str() {
            "succeeded" | "completed" => OperationProgress::Succeeded,
            "failed" => OperationProgress::Failed,
            "canceled" => OperationProgress::Canceled,
            _ => OperationProgress::Pending,
        };
        OperationObservation {
            progress,
            next_poll_after: operation.next_poll_after_ms.map(Duration::from_millis),
            error_code: operation.error_code.clone(),
        }
    }
}