superdoc-sdk 0.1.0-next.beta.1

use crate::{
    generated::{
        operations::DocApi,
        types::{DocCloseBoundParams, DocCloseResult, DocOpenParams, DocOpenResult},
    },
    invoke::Invoker,
    runtime::{
        contract, policy,
        process::{CallState, HostMode, HostOptions, Transport},
        routing,
    },
    Error, InvokeOptions, Result,
};
use base64::{engine::general_purpose::STANDARD as BASE64, Engine};
use serde_json::{json, Value};
use std::{
    collections::BTreeSet,
    path::Path,
    sync::{
        atomic::{AtomicU64, AtomicU8, Ordering},
        Arc, Mutex,
    },
    time::{Duration, SystemTime, UNIX_EPOCH},
};
use tokio::sync::broadcast;

#[derive(Clone)]
pub struct CollaborationAuth {
    token: String,
}
impl CollaborationAuth {
    pub fn token(token: impl Into<String>) -> Result<Self> {
        let token = token.into();
        if token.is_empty() {
            return Err(Error::new(
                "INVALID_ARGUMENT",
                "collaboration token must not be empty",
            ));
        }
        Ok(Self { token })
    }
    fn json(&self) -> Value {
        json!({"type":"token","token":self.token})
    }
}
impl std::fmt::Debug for CollaborationAuth {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str("CollaborationAuth([redacted])")
    }
}

#[derive(Clone, Debug, Default)]
pub struct OpenOptions {
    pub invoke: InvokeOptions,
    pub collaboration_auth: Option<CollaborationAuth>,
}

#[derive(Clone, Default)]
pub struct ClientOptions {
    pub default_change_mode: Option<String>,
    pub user_name: Option<String>,
    pub user_email: Option<String>,
    /// Opt in to negotiated document RPC on the CLI host. Collaboration and
    /// other ineligible opens keep their CLI route before any request is sent.
    pub document_rpc: bool,
}

struct ClientInner {
    transport: Transport,
    options: ClientOptions,
    opening: Mutex<BTreeSet<String>>,
}

#[derive(Clone)]
pub struct Client(Arc<ClientInner>);

struct OpenGuard {
    client: Client,
    id: String,
}
impl Drop for OpenGuard {
    fn drop(&mut self) {
        self.client
            .0
            .opening
            .lock()
            .expect("SDK opening lock")
            .remove(&self.id);
    }
}

impl Client {
    pub async fn connect(host: HostOptions) -> Result<Self> {
        Self::connect_with_options(host, ClientOptions::default()).await
    }

    pub async fn connect_with_options(host: HostOptions, options: ClientOptions) -> Result<Self> {
        if options
            .default_change_mode
            .as_deref()
            .is_some_and(|mode| !["direct", "tracked"].contains(&mode))
        {
            return Err(Error::new(
                "INVALID_ARGUMENT",
                "default change mode must be direct or tracked",
            ));
        }
        if options
            .user_name
            .as_deref()
            .is_some_and(|name| name.trim().is_empty())
            || (options.user_email.is_some() && options.user_name.is_none())
        {
            return Err(Error::new(
                "INVALID_ARGUMENT",
                "user identity requires a non-empty name",
            ));
        }
        Ok(Self(Arc::new(ClientInner {
            transport: Transport::connect(host).await?,
            options,
            opening: Mutex::new(BTreeSet::new()),
        })))
    }

    /// Local contract introspection requires no running host.
    pub fn describe() -> Value {
        contract()["introspection"]["overview"].clone()
    }

    pub fn describe_command(lookup: &str) -> Result<Value> {
        let index = contract()["introspection"]["lookup"][lookup.trim().to_lowercase()]
            .as_u64()
            .ok_or_else(|| {
                Error::new(
                    "OPERATION_NOT_FOUND",
                    "no contract operation matches the requested lookup",
                )
            })?;
        Ok(contract()["introspection"]["overview"]["operations"][index as usize].clone())
    }

    pub fn subscribe(&self) -> Notifications {
        Notifications(self.0.transport.subscribe())
    }
    pub fn state_directory(&self) -> &Path {
        self.0.transport.state_directory()
    }

    pub async fn open(&self, params: &DocOpenParams) -> Result<Document> {
        self.open_with_options(params, OpenOptions::default()).await
    }

    pub async fn open_with_options(
        &self,
        params: &DocOpenParams,
        options: OpenOptions,
    ) -> Result<Document> {
        let mut params = serde_json::to_value(params)
            .map_err(|_| Error::new("INVALID_ARGUMENT", "open parameters could not be encoded"))?;
        let session_id = match params.get("sessionId") {
            Some(value) => value
                .as_str()
                .filter(|id| !id.is_empty())
                .ok_or_else(|| Error::new("INVALID_ARGUMENT", "session ID must be non-empty"))?
                .to_owned(),
            None => fresh_session_id(),
        };
        {
            let mut opening = self.0.opening.lock().expect("SDK opening lock");
            if self.0.transport.has_session(&session_id) || !opening.insert(session_id.clone()) {
                return Err(Error::new(
                    "SESSION_ALREADY_OPEN",
                    "session ID is already open or opening",
                ));
            }
        }
        let _guard = OpenGuard {
            client: self.clone(),
            id: session_id.clone(),
        };
        params["sessionId"] = json!(session_id);
        let features = self.0.transport.features();
        let structured = options.collaboration_auth.is_none()
            && routing::can_open_document(&params, &options.invoke)
            && (self.0.transport.options().mode == HostMode::Document
                || self.0.options.document_rpc)
            && policy()["documentRpcFeatures"]
                .as_array()
                .is_some_and(|required| {
                    required
                        .iter()
                        .all(|feature| features.contains(feature.as_str().unwrap_or("")))
                });
        if self.0.transport.options().mode == HostMode::Document && !structured {
            return Err(Error::new(
                "DOCUMENT_RPC_INPUT_UNSUPPORTED",
                "standalone document host cannot support these open parameters",
            ));
        }
        let result = if structured {
            if self.0.options.default_change_mode.is_some() {
                self.require_feature("defaultChangeMode")?;
            }
            if self.0.options.user_name.is_some() {
                self.require_feature("documentAuthor")?;
            }
            let mut request = json!({"sessionId":session_id,"path":params["doc"]});
            if let Some(mode) = &self.0.options.default_change_mode {
                request["defaultChangeMode"] = json!(mode);
            }
            if let Some(name) = &self.0.options.user_name {
                let mut author = json!({"name":name});
                if let Some(email) = &self.0.options.user_email {
                    author["email"] = json!(email);
                }
                request["author"] = author;
            }
            self.rpc("document.open", request, &options.invoke, true, true, true)
                .await
                .and_then(|response| {
                    routing::map_lifecycle(
                        "doc.open",
                        response,
                        &session_id,
                        params["doc"].as_str(),
                    )
                })
        } else {
            if let Some(name) = &self.0.options.user_name {
                if params.get("userName").is_none() {
                    params["userName"] = json!(name);
                }
            }
            if let Some(email) = &self.0.options.user_email {
                if params.get("userEmail").is_none() {
                    params["userEmail"] = json!(email);
                }
            }
            self.invoke_cli(
                "doc.open",
                params,
                &options.invoke,
                options.collaboration_auth.as_ref(),
            )
            .await
        };
        let result = match result {
            Ok(value) => serde_json::from_value::<DocOpenResult>(value).map_err(Error::decode),
            Err(mut error) => {
                error.operation_id = Some("doc.open".to_owned());
                if ["TIMEOUT", "HOST_DISCONNECTED", "HOST_PROTOCOL_ERROR"]
                    .contains(&error.code.as_str())
                {
                    let _ = self.shutdown().await;
                    error.outcome_unknown = true;
                }
                return Err(error);
            }
        };
        let result = match result {
            Ok(result) if result.context_id == session_id && result.active => result,
            invalid => {
                // An invalid open result cannot leave an unreachable native
                // session alive. Retiring also invalidates any sibling handles.
                let _ = self.shutdown().await;
                let mut error = match invalid {
                    Err(error) => error,
                    Ok(_) => Error::new(
                        "HOST_PROTOCOL_ERROR",
                        "host returned an invalid open result",
                    ),
                };
                error.operation_id = Some("doc.open".into());
                error.outcome_unknown = true;
                return Err(error);
            }
        };
        let state = self.0.transport.session_state(&session_id);
        if state.load(Ordering::Acquire) != 0 {
            return Err(Error::new(
                "DOCUMENT_CLOSED",
                "session closed while it was opening",
            ));
        }
        Ok(Document(Arc::new(DocumentInner {
            client: self.clone(),
            session_id,
            state,
            structured,
            open_result: result,
        })))
    }

    pub async fn shutdown(&self) -> Result<()> {
        self.0.transport.shutdown().await
    }

    fn require_feature(&self, key: &str) -> Result<()> {
        if !self
            .0
            .transport
            .features()
            .contains(policy()["features"][key].as_str().unwrap_or(""))
        {
            return Err(Error::new(
                "CAPABILITY_UNSUPPORTED",
                "connected host lacks the requested optional capability",
            ));
        }
        Ok(())
    }

    async fn rpc(
        &self,
        method: &str,
        params: Value,
        options: &InvokeOptions,
        mutates: bool,
        open: bool,
        supports_deadline: bool,
    ) -> Result<Value> {
        self.rpc_tracked(
            method,
            params,
            options,
            mutates,
            open,
            supports_deadline,
            None,
        )
        .await
    }

    #[allow(clippy::too_many_arguments)]
    async fn rpc_tracked(
        &self,
        method: &str,
        params: Value,
        options: &InvokeOptions,
        mutates: bool,
        open: bool,
        supports_deadline: bool,
        state: Option<Arc<CallState>>,
    ) -> Result<Value> {
        let effective = options
            .timeout
            .unwrap_or(self.0.transport.options().request_timeout);
        let headroom =
            Duration::from_millis(policy()["watchdogHeadroomMs"].as_u64().unwrap_or(5000));
        if options.timeout.is_some() {
            self.require_feature("requestTimeout")?;
            if !supports_deadline {
                return Err(Error::new(
                    "DOCUMENT_RPC_INPUT_UNSUPPORTED",
                    "this native operation does not support an explicit request deadline",
                ));
            }
        }
        let deadline = (supports_deadline
            && self.0.transport.features().contains(
                policy()["features"]["requestTimeout"]
                    .as_str()
                    .unwrap_or(""),
            ))
        .then_some(effective);
        self.0
            .transport
            .request_tracked(
                method,
                params,
                deadline,
                effective.saturating_add(headroom),
                mutates,
                open,
                state.unwrap_or_default(),
            )
            .await
    }

    async fn invoke_cli(
        &self,
        id: &str,
        params: Value,
        options: &InvokeOptions,
        auth: Option<&CollaborationAuth>,
    ) -> Result<Value> {
        self.invoke_cli_tracked(id, params, options, auth, None)
            .await
    }

    async fn invoke_cli_tracked(
        &self,
        id: &str,
        params: Value,
        options: &InvokeOptions,
        auth: Option<&CollaborationAuth>,
        state: Option<Arc<CallState>>,
    ) -> Result<Value> {
        if self.0.transport.options().mode != HostMode::Cli {
            return Err(Error::new(
                "DOCUMENT_RPC_INPUT_UNSUPPORTED",
                "operation requires the CLI host",
            ));
        }
        let operation = operation(id)?;
        if auth.is_some() {
            if id != "doc.open" {
                return Err(Error::new(
                    "INVALID_ARGUMENT",
                    "authentication is scoped to document open",
                ));
            }
            self.require_feature("collaborationAuthPerOpen")?;
        }
        if options
            .stdin_bytes
            .as_ref()
            .is_some_and(|bytes| bytes.len() > 32 * 1024 * 1024)
        {
            return Err(Error::new(
                "PAYLOAD_TOO_LARGE",
                "stdin exceeds the host's 32 MiB limit",
            ));
        }
        validate_json(&params)?;
        let mut request = json!({"argv":routing::argv(operation,&params,options,self.0.options.default_change_mode.as_deref())?,"stdinBase64":options.stdin_bytes.as_ref().map(|bytes|BASE64.encode(bytes)).unwrap_or_default()});
        if let Some(auth) = auth {
            request["collaborationAuth"] = auth.json();
        }
        let response = self
            .rpc_tracked(
                "cli.invoke",
                request,
                options,
                operation["mutates"] == true,
                id == "doc.open",
                true,
                state,
            )
            .await?;
        routing::map_cli(operation, response).map_err(|mut error| {
            error.outcome_unknown = operation["mutates"] == true;
            error
        })
    }

    pub async fn get_tools(&self, provider: &str) -> Result<Value> {
        self.invoke_cli(
            "doc.preset.getTools",
            json!({"provider":provider,"preset":"core"}),
            &InvokeOptions::default(),
            None,
        )
        .await
    }

    pub async fn get_system_prompt(&self) -> Result<Value> {
        self.invoke_cli(
            "doc.preset.getSystemPrompt",
            json!({"preset":"core"}),
            &InvokeOptions::default(),
            None,
        )
        .await
    }
}

fn operation(id: &str) -> Result<&'static Value> {
    contract()["operations"].get(id).ok_or_else(|| {
        Error::new(
            "OPERATION_NOT_FOUND",
            "operation is not in the generated contract",
        )
    })
}

fn fresh_session_id() -> String {
    static NEXT: AtomicU64 = AtomicU64::new(1);
    format!(
        "rust-{}-{}-{}",
        std::process::id(),
        SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap_or_default()
            .as_nanos(),
        NEXT.fetch_add(1, Ordering::Relaxed)
    )
}

fn validate_json(value: &Value) -> Result<()> {
    let mut stack = vec![(value, 0)];
    while let Some((value, depth)) = stack.pop() {
        if depth > 128 {
            return Err(Error::new(
                "INVALID_ARGUMENT",
                "request JSON exceeds the supported nesting depth",
            ));
        }
        match value {
            Value::Number(number)
                if number.is_i64()
                    && number.as_i64().is_some_and(|n| {
                        !(-crate::wire::Integer::MAX..=crate::wire::Integer::MAX).contains(&n)
                    }) =>
            {
                return Err(Error::new(
                    "INVALID_ARGUMENT",
                    "JSON integer cannot be represented exactly by the host",
                ))
            }
            Value::Number(number)
                if number.is_u64()
                    && number
                        .as_u64()
                        .is_some_and(|n| n > crate::wire::Integer::MAX as u64) =>
            {
                return Err(Error::new(
                    "INVALID_ARGUMENT",
                    "JSON integer cannot be represented exactly by the host",
                ))
            }
            Value::Array(values) => stack.extend(values.iter().map(|value| (value, depth + 1))),
            Value::Object(values) => stack.extend(values.values().map(|value| (value, depth + 1))),
            _ => {}
        }
    }
    Ok(())
}

struct DocumentInner {
    client: Client,
    session_id: String,
    // Active, closing, closed, or outcome unknown. All clones share this state.
    state: Arc<AtomicU8>,
    structured: bool,
    open_result: DocOpenResult,
}

impl Drop for DocumentInner {
    fn drop(&mut self) {
        if self.state.load(Ordering::Acquire) >= 2 {
            self.client
                .0
                .transport
                .remove_session(&self.session_id, &self.state);
        }
    }
}

#[derive(Clone)]
pub struct Document(Arc<DocumentInner>);

struct CloseGuard {
    state: Arc<AtomicU8>,
    transmission: Arc<CallState>,
}
impl Drop for CloseGuard {
    fn drop(&mut self) {
        let _ = self.state.compare_exchange(
            1,
            if self.transmission.cancel_before_send() {
                0
            } else {
                3
            },
            Ordering::AcqRel,
            Ordering::Acquire,
        );
    }
}

impl Document {
    pub fn session_id(&self) -> &str {
        &self.0.session_id
    }
    pub fn open_result(&self) -> &DocOpenResult {
        &self.0.open_result
    }
    pub fn api(&self) -> DocApi<'_, Self> {
        DocApi::new(self)
    }

    pub async fn close(&self, params: &DocCloseBoundParams) -> Result<DocCloseResult> {
        self.api().close(params).await
    }

    pub async fn replace_file(
        &self,
        path: impl AsRef<Path>,
        options: InvokeOptions,
    ) -> Result<Value> {
        self.replace_path(path.as_ref(), options, None).await
    }

    /// The transport owns staged bytes until the native reply or confirmed exit,
    /// including when the caller drops this future after transmission.
    pub async fn replace_bytes(&self, bytes: &[u8], options: InvokeOptions) -> Result<Value> {
        use std::io::Write;
        if bytes.len() > 32 * 1024 * 1024 {
            return Err(Error::new(
                "PAYLOAD_TOO_LARGE",
                "replacement exceeds 32 MiB",
            ));
        }
        self.active()?;
        let mut file = tempfile::NamedTempFile::new_in(self.0.client.state_directory())
            .map_err(|_| Error::new("IO_ERROR", "could not stage replacement bytes"))?;
        file.write_all(bytes)
            .map_err(|_| Error::new("IO_ERROR", "could not stage replacement bytes"))?;
        let asset = file.into_temp_path();
        let path = asset.to_path_buf();
        self.replace_path(&path, options, Some(CallState::with_asset(asset)))
            .await
    }

    async fn replace_path(
        &self,
        path: &Path,
        options: InvokeOptions,
        state: Option<Arc<CallState>>,
    ) -> Result<Value> {
        self.active()?;
        self.0.client.require_feature("replaceFile")?;
        let path = path
            .to_str()
            .ok_or_else(|| Error::new("INVALID_ARGUMENT", "replacement path must be UTF-8"))?;
        let value = self
            .0
            .client
            .rpc_tracked(
                "document.replaceFile",
                json!({"sessionId":self.0.session_id,"path":path}),
                &options,
                true,
                false,
                false,
                state,
            )
            .await?;
        if value["sessionId"] != self.0.session_id
            || value["replaced"] != true
            || !value["byteLength"].is_number()
            || !value["revision"].is_string()
        {
            let mut error = Error::new(
                "HOST_PROTOCOL_ERROR",
                "host returned an invalid replacement result",
            );
            error.outcome_unknown = true;
            return Err(error);
        }
        Ok(
            json!({"contextId":self.0.session_id,"runtime":"v2","replaced":true,"document":{"byteLength":value["byteLength"],"revision":value["revision"]}}),
        )
    }

    pub async fn dispatch_tool(
        &self,
        name: &str,
        args: Value,
        options: InvokeOptions,
    ) -> Result<Value> {
        self.active()?;
        self.0
            .client
            .invoke_cli(
                "doc.preset.dispatch",
                json!({"sessionId":self.0.session_id,"toolName":name,"args":args,"preset":"core"}),
                &options,
                None,
            )
            .await
    }

    fn active(&self) -> Result<()> {
        match self.0.state.load(Ordering::Acquire) {
            0 if !self.0.client.0.transport.is_closed() => Ok(()),
            3 => Err(Error::new(
                "DOCUMENT_STATE_UNKNOWN",
                "document state is uncertain; use a new client to reopen verified saved content",
            )),
            _ => Err(Error::new(
                "DOCUMENT_CLOSED",
                "document handle is closed or closing",
            )),
        }
    }
}

impl Invoker for Document {
    async fn invoke(
        &self,
        id: &str,
        mut params: Value,
        mut options: InvokeOptions,
    ) -> Result<Value> {
        self.active()?;
        let operation = operation(id)?;
        if operation["sdkSurface"] != "document" {
            return Err(Error::new(
                "INVALID_ARGUMENT",
                "operation is not on the public document surface",
            ));
        }
        let binds_session = operation["params"]
            .as_array()
            .is_some_and(|params| params.iter().any(|param| param["name"] == "sessionId"));
        if binds_session {
            params["sessionId"] = json!(self.0.session_id);
        }
        let closing = id == "doc.close";
        let transmission = closing.then(|| Arc::new(CallState::default()));
        let _close = if closing {
            self.0
                .state
                .compare_exchange(0, 1, Ordering::AcqRel, Ordering::Acquire)
                .map_err(|_| Error::new("DOCUMENT_CLOSED", "document is already closing"))?;
            Some(CloseGuard {
                state: self.0.state.clone(),
                transmission: transmission.clone().expect("close transmission state"),
            })
        } else {
            None
        };
        if options.change_mode.is_none()
            && operation["params"]
                .as_array()
                .is_some_and(|params| params.iter().any(|param| param["name"] == "changeMode"))
        {
            options.change_mode = self.0.client.0.options.default_change_mode.clone();
        }
        let result: Result<Value> = async {
            if self.0.structured && binds_session {
                let (method, request) = routing::document_request(
                    &self.0.session_id,
                    operation,
                    &params,
                    &options,
                    self.0.client.0.transport.features(),
                )?;
                let safe_deadline = method != "document.save"
                    && (operation["mutates"] == false
                        || operation["idempotency"] == "idempotent"
                        || request["options"]["dryRun"] == true
                        || method != "document.invoke");
                validate_json(&request)?;
                let response = self
                    .0
                    .client
                    .rpc_tracked(
                        &method,
                        request,
                        &options,
                        operation["mutates"] == true,
                        false,
                        safe_deadline,
                        transmission.clone(),
                    )
                    .await;
                response.and_then(|response| {
                    if id == "doc.save" || closing {
                        routing::map_lifecycle(id, response, &self.0.session_id, None).map_err(
                            |mut error| {
                                error.outcome_unknown = operation["mutates"] == true;
                                error
                            },
                        )
                    } else {
                        Ok(response)
                    }
                })
            } else {
                self.0
                    .client
                    .invoke_cli_tracked(id, params, &options, None, transmission)
                    .await
            }
        }
        .await;
        let result = result.and_then(|value| routing::map_domain_result(operation, value));
        match result {
            Ok(value) => {
                if closing {
                    if value["closed"] != true || value["contextId"] != self.0.session_id {
                        let mut error = Error::new(
                            "HOST_PROTOCOL_ERROR",
                            "close was not confirmed by the host",
                        );
                        error.outcome_unknown = true;
                        return Err(error);
                    }
                    self.0.state.store(2, Ordering::Release);
                    self.0
                        .client
                        .0
                        .transport
                        .remove_session(&self.0.session_id, &self.0.state);
                }
                Ok(value)
            }
            Err(mut error) => {
                error.operation_id = Some(id.to_owned());
                if closing && !error.outcome_unknown && !self.0.client.0.transport.is_closed() {
                    let _ =
                        self.0
                            .state
                            .compare_exchange(1, 0, Ordering::AcqRel, Ordering::Acquire);
                }
                Err(error)
            }
        }
    }
}

/// Lag is explicit; callers must resynchronize rather than assume every remote
/// change notification was observed. Session invalidation runs in the reader's
/// owner independently of this public notification stream.
pub struct Notifications(broadcast::Receiver<Value>);
impl Notifications {
    pub async fn recv(&mut self) -> Result<Value> {
        self.0.recv().await.map_err(|error| match error {
            broadcast::error::RecvError::Lagged(_) => Error::new(
                "NOTIFICATIONS_LAGGED",
                "resynchronize document state before using notification-derived caches",
            ),
            broadcast::error::RecvError::Closed => {
                Error::new("HOST_DISCONNECTED", "notification stream closed")
            }
        })
    }
}