import { PublicKey } from "@solana/web3.js"
import { Tool } from "langchain/tools"

export class SolanaBalanceTool extends Tool {
  name = "solana_balance"
  description = `Get the balance of a Solana wallet or token account.

  If you want to get the balance of your wallet, you don't need to provide the tokenAddress.
  If no tokenAddress is provided, the balance will be in SOL.

  Inputs ( input is a JSON string ):
  tokenAddress: string, eg "So11111111111111111111111111111111111111112" (optional)`

  constructor(solanaKit) {
    super()
    this.solanaKit = solanaKit
  }

  async _call(input) {
    try {
      const tokenAddress = input ? new PublicKey(input) : undefined
      const balance = await this.solanaKit.getBalance(tokenAddress)

      return JSON.stringify({
        status: "success",
        balance,
        token: input || "SOL"
      })
    } catch (error) {
      return JSON.stringify({
        status: "error",
        message: error.message,
        code: error.code || "UNKNOWN_ERROR"
      })
    }
  }
}