From 8b36aca8b6bd054b13f62b83982e0e5ae1381108 Mon Sep 17 00:00:00 2001
From: callebtc <93376500+callebtc@users.noreply.github.com>
Date: Sat, 8 Oct 2022 14:27:41 +0200
Subject: [PATCH] keyset balances

---
 cashu/mint/ledger.py   |  2 +-
 cashu/wallet/cli.py    | 13 +++++++++++--
 cashu/wallet/wallet.py | 26 +++++++++++++++++++++-----
 3 files changed, 33 insertions(+), 8 deletions(-)

diff --git a/cashu/mint/ledger.py b/cashu/mint/ledger.py
index e3ab10c..d98553a 100644
--- a/cashu/mint/ledger.py
+++ b/cashu/mint/ledger.py
@@ -207,7 +207,7 @@ class Ledger:
 
     # Public methods
     def get_keyset(self):
-        return {"id": self.keyset.id, "keys": self._serialize_pubkeys()}
+        return self._serialize_pubkeys()
 
     async def request_mint(self, amount):
         """Returns Lightning invoice and stores it in the db."""
diff --git a/cashu/wallet/cli.py b/cashu/wallet/cli.py
index b45ee31..7c7aeb5 100755
--- a/cashu/wallet/cli.py
+++ b/cashu/wallet/cli.py
@@ -148,7 +148,16 @@ async def pay(ctx, invoice: str):
 @coro
 async def balance(ctx):
     wallet: Wallet = ctx.obj["WALLET"]
-    wallet.status()
+    keyset_balances = wallet.balance_per_keyset()
+    if len(keyset_balances) > 1:
+        for k, v in keyset_balances.items():
+            print(
+                f"Keyset: {k} Balance: {v['balance']} sat (available: {v['available']})"
+            )
+        print("")
+    print(
+        f"Balance: {wallet.balance} sat (available: {wallet.available_balance} sat in {len([p for p in wallet.proofs if not p.reserved])} tokens)"
+    )
 
 
 @cli.command("send", help="Send coins.")
@@ -192,7 +201,7 @@ async def receive(ctx, coin: str, lock: str):
         address_split = lock.split("P2SH:")[1]
 
         p2shscripts = await get_unused_locks(address_split, db=wallet.db)
-        assert len(p2shscripts) == 1
+        assert len(p2shscripts) == 1, Exception("lock not found.")
         script = p2shscripts[0].script
         signature = p2shscripts[0].signature
     else:
diff --git a/cashu/wallet/wallet.py b/cashu/wallet/wallet.py
index 5a42241..4d9f6d9 100644
--- a/cashu/wallet/wallet.py
+++ b/cashu/wallet/wallet.py
@@ -3,6 +3,7 @@ import json
 import secrets as scrts
 import uuid
 from typing import Dict, List
+from itertools import groupby
 
 import requests
 from loguru import logger
@@ -50,15 +51,13 @@ class LedgerAPI:
 
     def _get_keys(self, url):
         resp = requests.get(url + "/keys").json()
-        keyset_id = resp["id"]
-        keys = resp["keys"]
+        keys = resp
         assert len(keys), Exception("did not receive any keys")
         keyset_keys = {
             int(amt): PublicKey(bytes.fromhex(val), raw=True)
             for amt, val in keys.items()
         }
         keyset = Keyset(pubkeys=keyset_keys)
-        assert keyset_id == keyset.id, Exception("mint keyset id not valid.")
         return keyset
 
     @staticmethod
@@ -358,7 +357,6 @@ class Wallet(LedgerAPI):
         return token
 
     async def _get_spendable_proofs(self, proofs: List[Proof]):
-        print(f"Debug: only loading proofs with id: {self.keyset_id}")
         proofs = [p for p in proofs if p.id == self.keyset_id or not p.id]
         proofs = [p for p in proofs if not p.reserved]
         return proofs
@@ -368,7 +366,6 @@ class Wallet(LedgerAPI):
         if scnd_secret:
             logger.debug(f"Spending conditions: {scnd_secret}")
         spendable_proofs = await self._get_spendable_proofs(proofs)
-        print(f"Balance: {sum([p.amount for p in spendable_proofs])}")
         if sum([p.amount for p in spendable_proofs]) < amount:
             raise Exception("balance too low.")
         return await self.split(
@@ -428,5 +425,24 @@ class Wallet(LedgerAPI):
             f"Balance: {self.balance} sat (available: {self.available_balance} sat in {len([p for p in self.proofs if not p.reserved])} tokens)"
         )
 
+    @staticmethod
+    def _sum_proofs(proofs: List[Proof], available_only=False):
+        if available_only:
+            return sum([p.amount for p in proofs if not p.reserved])
+        return sum([p.amount for p in proofs])
+
+    @staticmethod
+    def _get_proofs_per_keyset(proofs: List[Proof]):
+        return {key: list(group) for key, group in groupby(proofs, lambda p: p.id)}
+
+    def balance_per_keyset(self):
+        return {
+            key: {
+                "balance": self._sum_proofs(proofs),
+                "available": self._sum_proofs(proofs, available_only=True),
+            }
+            for key, proofs in self._get_proofs_per_keyset(self.proofs).items()
+        }
+
     def proof_amounts(self):
         return [p["amount"] for p in sorted(self.proofs, key=lambda p: p["amount"])]