use crate::{
    commands::execute_command,
    llms::tools::{ToolDefinition, ToolExecution},
};
use schemars::{schema::RootSchema, schema_for, JsonSchema};
use serde::{Deserialize, Serialize};

use super::TypedTool;
use anyhow::Result;

#[derive(Debug, Serialize, Deserialize, JsonSchema)]
pub struct ProgramArgs {
    program_name: String,
    program_args: Option<Vec<String>>,
}

pub struct AnyProgramTool;

impl TypedTool for AnyProgramTool {
    type Args = ProgramArgs;
    type Output = String;

    fn name(&self) -> &str {
        "any-program"
    }

    fn description(&self) -> &str {
        "Execute any program available to the user on the local system"
    }

    fn execute(&self, args: Self::Args) -> Result<Self::Output> {
        let args_vec: Vec<&str> = args
            .program_args
            .as_ref()
            .map(|args| args.iter().map(|s| s.as_str()).collect())
            .unwrap_or_default();

        Ok(execute_command(
            &args.program_name,
            Some(args_vec.as_slice()),
        )?)
    }
}

#[derive(Debug, Serialize, Deserialize, JsonSchema)]
struct LocalProgram {
    program_name: String,
    description: String,
    program_args: Option<Vec<String>>,
}

impl ToolDefinition for LocalProgram {
    fn name(&self) -> String {
        self.program_name.clone()
    }

    fn description(&self) -> String {
        self.description.clone()
    }

    fn args(&self) -> RootSchema {
        schema_for!(Self)
    }
}

impl ToolExecution for LocalProgram {
    type Args = Vec<String>;
    type Return = String;

    fn schema(&self) -> RootSchema {
        self.args()
    }

    fn args_from_schema(&self, schema: serde_json::Value) -> anyhow::Result<Self::Args> {
        Ok(schema
            .as_array()
            .unwrap()
            .iter()
            .map(|v| v.as_str().unwrap().to_string())
            .collect())
    }

    fn json_schema_to_args(&self, schema: RootSchema) -> anyhow::Result<Self::Args> {
        self.args_from_schema(serde_json::to_value(schema)?)
    }

    fn args_from_json(&self, json: serde_json::Value) -> anyhow::Result<Self::Args> {
        self.args_from_schema(json)
    }

    fn execute(&self, args: Self::Args) -> anyhow::Result<Self::Return> {
        let str_args: Vec<&str> = args.iter().map(|s| s.as_str()).collect();
        Ok(execute_command(
            &self.program_name,
            Some(str_args.as_slice()),
        )?)
    }
}
