//! End-to-end example:
//!
//! - Register two tiny tools (`add` = Read, `remember_project` = Write) into a
//!   `ToolRegistry`.
//! - Spin up an MCP server on a random localhost port so external clients
//!   (Claude Desktop, etc.) could connect.
//! - Drive the same registry in-process via an `Agent<OllamaProvider>`, with
//!   the `remember_project` capability granted to the session.
//!
//! Run with a tool-calling model installed:
//!
//! ```
//! ollama pull qwen2.5
//! cargo run --example toy
//! ```
//!
//! Override the model with `KBERG_MODEL=llama3.2 cargo run --example toy`.

use std::collections::HashSet;

use async_trait::async_trait;
use kberg::{
    Agent, Error, Tool, ToolCallResult, ToolKind, ToolRegistry, WriteCapability,
    provider::ollama::OllamaProvider, server,
};
use serde_json::{Value, json};

struct Add;

#[async_trait]
impl Tool for Add {
    fn name(&self) -> &'static str {
        "add"
    }
    fn description(&self) -> &'static str {
        "Add two integers and return the sum."
    }
    fn kind(&self) -> ToolKind {
        ToolKind::Read
    }
    fn small_model_safe(&self) -> bool {
        true
    }
    fn input_schema(&self) -> Value {
        json!({
            "type": "object",
            "properties": {
                "a": { "type": "integer" },
                "b": { "type": "integer" }
            },
            "required": ["a", "b"]
        })
    }
    async fn call(&self, args: Value) -> Result<ToolCallResult, Error> {
        let a = args
            .get("a")
            .and_then(Value::as_i64)
            .ok_or_else(|| Error::InvalidArgs {
                tool: "add".into(),
                message: "missing integer `a`".into(),
            })?;
        let b = args
            .get("b")
            .and_then(Value::as_i64)
            .ok_or_else(|| Error::InvalidArgs {
                tool: "add".into(),
                message: "missing integer `b`".into(),
            })?;
        Ok(ToolCallResult::text((a + b).to_string()))
    }
}

struct RememberProject;

#[async_trait]
impl Tool for RememberProject {
    fn name(&self) -> &'static str {
        "remember_project"
    }
    fn description(&self) -> &'static str {
        "Record a new project the user is working on."
    }
    fn kind(&self) -> ToolKind {
        ToolKind::Write(WriteCapability::new(
            "remember_project",
            "record new projects the user says they're working on",
        ))
    }
    fn input_schema(&self) -> Value {
        json!({
            "type": "object",
            "properties": { "name": { "type": "string" } },
            "required": ["name"]
        })
    }
    async fn call(&self, args: Value) -> Result<ToolCallResult, Error> {
        let name = args
            .get("name")
            .and_then(Value::as_str)
            .unwrap_or("<unnamed>");
        Ok(ToolCallResult::text(format!("recorded project `{name}`")))
    }
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    tracing_subscriber::fmt()
        .with_env_filter(
            tracing_subscriber::EnvFilter::try_from_default_env()
                .unwrap_or_else(|_| "info,kberg=debug".into()),
        )
        .init();

    let mut registry = ToolRegistry::new();
    registry.register(Add);
    registry.register(RememberProject);

    // Bring up an MCP server so external clients could reach the same tools.
    let bound = server::serve(
        registry.clone(),
        "127.0.0.1:0",
        server::ServeConfig {
            projection: kberg::SurfaceProjection::Full,
            grants: Some(HashSet::from(["remember_project".to_string()])),
            resources: None,
        },
    )
    .await?;
    println!("MCP server: {}", bound.url());
    println!(
        "registered write capabilities: {:?}",
        registry.write_capabilities()
    );

    let model = std::env::var("KBERG_MODEL").unwrap_or_else(|_| "qwen2.5".to_string());
    let provider = OllamaProvider::new(&model);
    let agent = Agent::new(provider).with_grants(HashSet::from(["remember_project".to_string()]));

    let prompt = "What is 17 + 25? Also, remember that I'm working on a project called Kberg.";

    println!("\n> {prompt}\n");
    let outcome = agent
        .run(
            &registry,
            Some(
                "You are a terse assistant. Prefer tools over guessing. \
                 When you have the answer, reply directly and briefly.",
            ),
            prompt,
        )
        .await?;

    println!(
        "--- final ({} step{}) ---",
        outcome.steps,
        if outcome.steps == 1 { "" } else { "s" }
    );
    println!("{}", outcome.final_message);

    bound.shutdown();
    Ok(())
}