procyon 0.1.2

Terminal development harness for Stellar and Soroban smart contracts, driven by a language model
use async_trait::async_trait;
use serde_json::{json, Value};

use super::Tool;
use crate::account::{get_accounts_path, query_balance, Account, AccountStore};
use crate::project::{Network, Project};

pub struct AccountCreateTool;

#[async_trait]
impl Tool for AccountCreateTool {
    fn name(&self) -> &str {
        "account_create"
    }

    fn capability(&self) -> crate::risk::Capability {
        crate::risk::Capability::Write
    }

    fn description(&self) -> &str {
        "Create a new Stellar account or import an existing one"
    }

    fn input_schema(&self) -> Value {
        json!({
            "type": "object",
            "properties": {
                "name": {
                    "type": "string",
                    "description": "Account name"
                },
                "address": {
                    "type": "string",
                    "description": "Stellar address (if importing existing account)"
                },
                "network": {
                    "type": "string",
                    "enum": ["local", "testnet", "mainnet"],
                    "description": "Network for this account (default: testnet)"
                }
            },
            "required": ["name"]
        })
    }

    async fn execute(&self, input: Value) -> Result<String, String> {
        let name = input
            .get("name")
            .and_then(|v| v.as_str())
            .ok_or("Missing 'name' parameter")?;

        let address = input
            .get("address")
            .and_then(|v| v.as_str())
            .map(|s| s.to_string());

        let network_str = input
            .get("network")
            .and_then(|v| v.as_str())
            .unwrap_or("testnet");

        let network: Network = network_str
            .parse()
            .map_err(|e: String| format!("Invalid network: {}", e))?;

        let current_dir = std::env::current_dir()
            .map_err(|e| format!("Failed to get current directory: {}", e))?;

        let project_dir = Project::find_project_dir(&current_dir)
            .await
            .map_err(|e| e.to_string())?;
        let accounts_path = get_accounts_path(&project_dir);
        let mut store = AccountStore::load(&accounts_path)
            .await
            .map_err(|e| e.to_string())?;

        if store.get(name).is_some() {
            return Err(format!("Account '{}' already exists", name));
        }

        let account_address = match address {
            Some(addr) => addr,
            None => {
                let generate = tokio::process::Command::new("stellar")
                    .args([
                        "keys",
                        "generate",
                        "--global",
                        name,
                        "--network",
                        network_str,
                        "--fund",
                    ])
                    .output()
                    .await
                    .map_err(|e| format!("Failed to generate key: {}", e))?;

                if !generate.status.success() {
                    let stderr = String::from_utf8_lossy(&generate.stderr);
                    return Err(format!("Failed to generate key: {}", stderr));
                }

                // `keys generate` does not print the public key, so it is read back explicitly.
                let show = tokio::process::Command::new("stellar")
                    .args(["keys", "address", name])
                    .output()
                    .await
                    .map_err(|e| format!("Failed to read generated address: {}", e))?;

                if !show.status.success() {
                    let stderr = String::from_utf8_lossy(&show.stderr);
                    return Err(format!("Failed to read generated address: {}", stderr));
                }

                String::from_utf8_lossy(&show.stdout).trim().to_string()
            }
        };

        let account = Account {
            name: name.to_string(),
            address: account_address.clone(),
            network,
        };

        store.add(account);
        store
            .save(&accounts_path)
            .await
            .map_err(|e| format!("Failed to save accounts: {}", e))?;

        Ok(format!(
            "Account created successfully!\n\nName: {}\nAddress: {}\nNetwork: {}",
            name, account_address, network_str
        ))
    }
}

pub struct AccountListTool;

#[async_trait]
impl Tool for AccountListTool {
    fn name(&self) -> &str {
        "account_list"
    }

    fn capability(&self) -> crate::risk::Capability {
        crate::risk::Capability::ReadOnly
    }

    fn description(&self) -> &str {
        "List all configured accounts"
    }

    fn input_schema(&self) -> Value {
        json!({
            "type": "object",
            "properties": {},
            "required": []
        })
    }

    async fn execute(&self, _input: Value) -> Result<String, String> {
        let current_dir = std::env::current_dir()
            .map_err(|e| format!("Failed to get current directory: {}", e))?;

        let project_dir = Project::find_project_dir(&current_dir)
            .await
            .map_err(|e| e.to_string())?;
        let accounts_path = get_accounts_path(&project_dir);
        let store = AccountStore::load(&accounts_path)
            .await
            .map_err(|e| e.to_string())?;

        let accounts = store.list();

        if accounts.is_empty() {
            return Ok("No accounts configured. Use account_create to add one.".to_string());
        }

        let mut output = format!("Accounts ({}):\n\n", accounts.len());
        for account in accounts {
            output.push_str(&format!(
                "  {} - {} [{}]\n",
                account.name, account.address, account.network
            ));
        }

        Ok(output)
    }
}

pub struct AccountBalanceTool;

#[async_trait]
impl Tool for AccountBalanceTool {
    fn name(&self) -> &str {
        "account_balance"
    }

    fn capability(&self) -> crate::risk::Capability {
        crate::risk::Capability::ReadOnly
    }

    fn description(&self) -> &str {
        "Query the balance of a Stellar account"
    }

    fn input_schema(&self) -> Value {
        json!({
            "type": "object",
            "properties": {
                "address": {
                    "type": "string",
                    "description": "Account address (or account name)"
                },
                "network": {
                    "type": "string",
                    "enum": ["local", "testnet", "mainnet"],
                    "description": "Network to query (default: testnet)"
                }
            },
            "required": ["address"]
        })
    }

    async fn execute(&self, input: Value) -> Result<String, String> {
        let address = input
            .get("address")
            .and_then(|v| v.as_str())
            .ok_or("Missing 'address' parameter")?;

        let network_str = input
            .get("network")
            .and_then(|v| v.as_str())
            .unwrap_or("testnet");

        let network: Network = network_str
            .parse()
            .map_err(|e: String| format!("Invalid network: {}", e))?;

        let resolved_address = resolve_account_name(address)
            .await
            .unwrap_or_else(|| address.to_string());

        query_balance(&resolved_address, &network)
            .await
            .map_err(|e| e.to_string())
    }
}

async fn resolve_account_name(name: &str) -> Option<String> {
    let current_dir = std::env::current_dir().ok()?;
    let project_dir = Project::find_project_dir(&current_dir).await.ok()?;
    let accounts_path = get_accounts_path(&project_dir);
    let store = AccountStore::load(&accounts_path).await.ok()?;
    store.get(name).map(|a| a.address.clone())
}