lodestone-client 0.0.4

use std::error::Error;
use std::fmt::{self, Display, Formatter};

use lodestone_wire::{DiscoverResult, LODESTONE_API_SCHEMA_ID, LODESTONE_API_SPEC_SHA256};
use serde_json::Value;

use crate::generated_tool_packs::{TOOL_INDEX_JSON, tool_pack_json};

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum ToolPackError {
    DiscoverySchemaMismatch,
    UnknownFamily,
    GeneratedArtifactInvalid,
}

impl Display for ToolPackError {
    fn fmt(&self, formatter: &mut Formatter<'_>) -> fmt::Result {
        formatter.write_str(match self {
            Self::DiscoverySchemaMismatch => "native discovery schema does not match the packed client",
            Self::UnknownFamily => "native tool capability family is unknown",
            Self::GeneratedArtifactInvalid => "generated native tool artifact is invalid",
        })
    }
}

impl Error for ToolPackError {}

/// Loads one generated capability-family pack after exact schema identity,
/// observed capability, and host result-ceiling intersection.
pub fn load_tool_pack(
    discovery: &DiscoverResult,
    family: &str,
    host_maximum_result_bytes: u64,
) -> Result<Value, ToolPackError> {
    if discovery.schema_id != LODESTONE_API_SCHEMA_ID || discovery.schema_digest.value != LODESTONE_API_SPEC_SHA256 {
        return Err(ToolPackError::DiscoverySchemaMismatch);
    }
    let index: Value = serde_json::from_str(TOOL_INDEX_JSON).map_err(|_| ToolPackError::GeneratedArtifactInvalid)?;
    if index.get("native_schema_sha256").and_then(Value::as_str) != Some(LODESTONE_API_SPEC_SHA256)
        || !index.get("families").and_then(Value::as_array).is_some_and(|families| {
            families
                .iter()
                .any(|entry| entry.get("id").and_then(Value::as_str) == Some(family))
        })
    {
        return Err(ToolPackError::UnknownFamily);
    }
    let source = tool_pack_json(family).ok_or(ToolPackError::UnknownFamily)?;
    let mut pack: Value = serde_json::from_str(source).map_err(|_| ToolPackError::GeneratedArtifactInvalid)?;
    if pack.get("native_schema_sha256").and_then(Value::as_str) != Some(LODESTONE_API_SPEC_SHA256) {
        return Err(ToolPackError::GeneratedArtifactInvalid);
    }
    let available = discovery
        .capabilities
        .iter()
        .filter(|capability| capability.state == "available" && capability.capability_family == family)
        .map(|capability| capability.command.as_str())
        .collect::<std::collections::BTreeSet<_>>();
    let operations = pack
        .get_mut("operations")
        .and_then(Value::as_array_mut)
        .ok_or(ToolPackError::GeneratedArtifactInvalid)?;
    operations.retain(|operation| {
        let command = operation.get("command").and_then(Value::as_str);
        let maximum = operation.get("maximum_result_bytes").and_then(Value::as_u64);
        command.is_some_and(|command| available.contains(command))
            && maximum.is_some_and(|maximum| maximum <= host_maximum_result_bytes)
    });
    Ok(pack)
}

/// Resolves one command's generated provider-neutral operation descriptor
/// without requiring a connection or inventing a second command registry.
pub fn describe_tool_operation(command: &str) -> Result<Option<Value>, ToolPackError> {
    let index: Value = serde_json::from_str(TOOL_INDEX_JSON).map_err(|_| ToolPackError::GeneratedArtifactInvalid)?;
    if index.get("native_schema_sha256").and_then(Value::as_str) != Some(LODESTONE_API_SPEC_SHA256) {
        return Err(ToolPackError::GeneratedArtifactInvalid);
    }
    let Some(family) = index.get("families").and_then(Value::as_array).and_then(|families| {
        families.iter().find(|family| {
            family
                .get("commands")
                .and_then(Value::as_array)
                .is_some_and(|commands| commands.iter().any(|value| value.as_str() == Some(command)))
        })
    }) else {
        return Ok(None);
    };
    let family = family
        .get("id")
        .and_then(Value::as_str)
        .ok_or(ToolPackError::GeneratedArtifactInvalid)?;
    let source = tool_pack_json(family).ok_or(ToolPackError::GeneratedArtifactInvalid)?;
    let pack: Value = serde_json::from_str(source).map_err(|_| ToolPackError::GeneratedArtifactInvalid)?;
    if pack.get("native_schema_sha256").and_then(Value::as_str) != Some(LODESTONE_API_SPEC_SHA256) {
        return Err(ToolPackError::GeneratedArtifactInvalid);
    }
    let mut operation = pack
        .get("operations")
        .and_then(Value::as_array)
        .and_then(|operations| {
            operations
                .iter()
                .find(|operation| operation.get("command").and_then(Value::as_str) == Some(command))
        })
        .cloned()
        .ok_or(ToolPackError::GeneratedArtifactInvalid)?;
    let definitions = pack
        .get("$defs")
        .cloned()
        .ok_or(ToolPackError::GeneratedArtifactInvalid)?;
    operation
        .as_object_mut()
        .ok_or(ToolPackError::GeneratedArtifactInvalid)?
        .insert("$defs".to_owned(), definitions);
    Ok(Some(operation))
}