chunk-your-tools 1.0.3

#![allow(clippy::multiple_crate_versions)]

use chunk_your_tools::{
    DecomposedCatalog, NamedSurvivors, PolicyContext, RetrieveOptions, build_catalog_from_tools,
    build_process_groups_options, parse_tool_policy, parse_tool_policy_pair,
    per_tool_policies_from_value, policy_context_from_values, recompose_tools_from_names,
    retrieve_tools_from_catalog,
};
use clap::{Parser, Subcommand};
use serde_json::Value;
use std::collections::HashMap;
use std::fs;
use std::path::{Path, PathBuf};

#[derive(Parser)]
#[command(name = "chunk-your-tools")]
#[command(about = "Decompose and recompose MCP tool definition schemas")]
struct Cli {
    #[command(subcommand)]
    command: Commands,
}

#[derive(Subcommand)]
enum Commands {
    /// Decompose MCP tools JSON into a catalog directory
    Decompose {
        #[arg(long)]
        input: PathBuf,
        #[arg(long)]
        output: PathBuf,
    },
    /// Recompose pruned tools from tools JSON and semantic survivor names
    Recompose {
        #[arg(long)]
        input: PathBuf,
        #[arg(long)]
        survivors: PathBuf,
        #[arg(long)]
        output: PathBuf,
        #[arg(long)]
        config: Option<PathBuf>,
        #[arg(long)]
        system_policy: Option<String>,
        #[arg(long)]
        mcp_policy: Option<String>,
        #[arg(long)]
        per_tool: Option<PathBuf>,
        #[arg(long = "tool-policy", value_name = "TOOL=POLICY")]
        tool_policies: Vec<String>,
    },
}

fn load_tools_array(tools: &Path) -> Result<Vec<Value>, Box<dyn std::error::Error>> {
    let raw = fs::read_to_string(tools)?;
    let tools_val: Value = serde_json::from_str(&raw)?;
    tools_val
        .as_array()
        .cloned()
        .or_else(|| tools_val.get("tools").and_then(|v| v.as_array()).cloned())
        .ok_or_else(|| "Expected tools array in input JSON".into())
}

fn policy_context_from_cli(
    config: Option<&Path>,
    system_policy: Option<&str>,
    mcp_policy: Option<&str>,
    per_tool: Option<&Path>,
    tool_policies: &[String],
) -> Result<PolicyContext, Box<dyn std::error::Error>> {
    let mut ctx = match config {
        Some(path) => {
            let raw = fs::read_to_string(path)?;
            policy_context_from_values(&serde_json::from_str(&raw)?)
        }
        None => PolicyContext::new(),
    };

    if let Some(s) = system_policy {
        ctx.system_policy =
            parse_tool_policy(s).ok_or_else(|| format!("invalid system policy: {s}"))?;
    }
    if let Some(m) = mcp_policy {
        ctx.mcp_policy = parse_tool_policy(m).ok_or_else(|| format!("invalid mcp policy: {m}"))?;
    }

    if let Some(path) = per_tool {
        let raw = fs::read_to_string(path)?;
        let val: Value = serde_json::from_str(&raw)?;
        chunk_your_tools::apply_per_tool_overrides(&mut ctx, per_tool_policies_from_value(&val)?);
    }

    let mut cli_overrides = HashMap::new();
    for spec in tool_policies {
        let (tool_id, policy) = parse_tool_policy_pair(spec)?;
        cli_overrides.insert(tool_id, policy);
    }
    chunk_your_tools::apply_per_tool_overrides(&mut ctx, cli_overrides);

    Ok(ctx)
}

fn run_decompose(input: &Path, output: &Path) -> Result<(), Box<dyn std::error::Error>> {
    let tools_arr = load_tools_array(input)?;
    let index = build_catalog_from_tools(&tools_arr);
    fs::create_dir_all(output)?;
    for (rel, content) in &index.files {
        let path = output.join(rel);
        if let Some(parent) = path.parent() {
            fs::create_dir_all(parent)?;
        }
        fs::write(path, content)?;
    }
    eprintln!("Wrote {} files to {}", index.files.len(), output.display());
    Ok(())
}

fn recompose_tools(
    tools_arr: &[Value],
    survivors_val: &Value,
    ctx: &PolicyContext,
) -> Result<Vec<Value>, Box<dyn std::error::Error>> {
    if survivors_val.get("tools").is_some() {
        let survivors = NamedSurvivors::from_value(survivors_val)?;
        return Ok(recompose_tools_from_names(tools_arr, &survivors, ctx));
    }
    if survivors_val.get("json").is_some() || survivors_val.get("md").is_some() {
        let index = build_catalog_from_tools(tools_arr);
        let build_catalog = index.to_catalog_dict();
        let mut store = DecomposedCatalog::from_catalog_index(&index);
        let process_groups = build_process_groups_options(ctx, &build_catalog, &store, None);
        let opts = RetrieveOptions {
            apply_decomposed_score_filter: false,
            process_groups,
        };
        return Ok(retrieve_tools_from_catalog(
            ctx,
            survivors_val,
            &build_catalog,
            &mut store,
            &opts,
        ));
    }
    Err(
        "survivors JSON must include \"tools\" (semantic names) or \"json\"/\"md\" (legacy chunks)"
            .into(),
    )
}

#[allow(clippy::too_many_arguments)]
fn run_recompose(
    input: &Path,
    survivors_path: &Path,
    output: &Path,
    config: Option<&Path>,
    system_policy: Option<&str>,
    mcp_policy: Option<&str>,
    per_tool: Option<&Path>,
    tool_policies: &[String],
) -> Result<(), Box<dyn std::error::Error>> {
    let ctx = policy_context_from_cli(config, system_policy, mcp_policy, per_tool, tool_policies)?;
    let tools_arr = load_tools_array(input)?;
    let survivors_raw = fs::read_to_string(survivors_path)?;
    let survivors_val: Value = serde_json::from_str(&survivors_raw)?;
    let tools = recompose_tools(&tools_arr, &survivors_val, &ctx)?;
    let mut json = serde_json::to_string_pretty(&tools)?;
    json.push('\n');
    if let Some(parent) = output.parent() {
        fs::create_dir_all(parent)?;
    }
    fs::write(output, json)?;
    eprintln!("Wrote {} tools to {}", tools.len(), output.display());
    Ok(())
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let cli = Cli::parse();
    match cli.command {
        Commands::Decompose { input, output } => run_decompose(&input, &output)?,
        Commands::Recompose {
            input,
            survivors,
            output,
            config,
            system_policy,
            mcp_policy,
            per_tool,
            tool_policies,
        } => run_recompose(
            &input,
            &survivors,
            &output,
            config.as_deref(),
            system_policy.as_deref(),
            mcp_policy.as_deref(),
            per_tool.as_deref(),
            &tool_policies,
        )?,
    }
    Ok(())
}