pub mod registry;
pub mod service;

use queryd::context::QueryEngine;
use arrow::json::writer::{JsonArray, Writer as JsonWriter};

/// State for the tool system.
#[derive(Clone)]
pub struct ToolState {
    pub registry: registry::ToolRegistry,
    pub query_fn: QueryExecutor,
}

/// Wraps QueryEngine to provide a simple execute interface for tools.
#[derive(Clone)]
pub struct QueryExecutor {
    engine: QueryEngine,
}

impl QueryExecutor {
    pub fn new(engine: QueryEngine) -> Self {
        Self { engine }
    }

    /// Execute SQL and return (rows as JSON, row count).
    pub async fn execute(&self, sql: &str) -> Result<(serde_json::Value, usize), String> {
        let batches = self.engine.query(sql).await?;

        if batches.is_empty() {
            return Ok((serde_json::Value::Array(vec![]), 0));
        }

        let mut buf = Vec::new();
        let mut writer = JsonWriter::<_, JsonArray>::new(&mut buf);
        for batch in &batches {
            writer.write(batch).map_err(|e| format!("JSON write: {e}"))?;
        }
        writer.finish().map_err(|e| format!("JSON finish: {e}"))?;
        drop(writer);

        let rows: serde_json::Value = serde_json::from_slice(&buf)
            .map_err(|e| format!("JSON parse: {e}"))?;
        let count = rows.as_array().map(|a| a.len()).unwrap_or(0);

        Ok((rows, count))
    }
}