From c31f312505e4a853a8cf81731fb8bcba43efdf68 Mon Sep 17 00:00:00 2001
From: Shroominic <db2002dominic@gmail.com>
Date: Thu, 5 Jun 2025 14:59:59 +0200
Subject: [PATCH] fix refund and improve wallet endpoints

---
 router/account.py | 60 +++++++++++++++++++++++++++++++++--------------
 1 file changed, 43 insertions(+), 17 deletions(-)

diff --git a/router/account.py b/router/account.py
index d7365806..9c90ba3f 100644
--- a/router/account.py
+++ b/router/account.py
@@ -1,12 +1,11 @@
 from typing import Annotated
 from fastapi import APIRouter, Header, HTTPException, Depends
-from sixty_nuts import Wallet
 
 from .auth import validate_bearer_key
-from .cashu import refund_balance, credit_balance, NSEC, MINT
+from .cashu import refund_balance, credit_balance, WALLET
 from .db import ApiKey, AsyncSession, get_session
 
-account_router = APIRouter(prefix="/v1/wallet")
+wallet_router = APIRouter(prefix="/v1/wallet")
 
 
 async def get_key_from_header(
@@ -21,17 +20,24 @@ async def get_key_from_header(
         detail="Invalid authorization. Use 'Bearer <cashu-token>' or 'Bearer <api-key>'",
     )
 
-
-@account_router.get("/")
+# TODO: remove this endpoint when frontend is updated
+@wallet_router.get("/")
 async def account_info(key: ApiKey = Depends(get_key_from_header)) -> dict:
     return {
         "api_key": "sk-" + key.hashed_key,
         "balance": key.balance,
     }
 
+@wallet_router.get("/info")
+async def wallet_info(key: ApiKey = Depends(get_key_from_header)) -> dict:
+    return {
+        "api_key": "sk-" + key.hashed_key,
+        "balance": key.balance,
+    }
 
-@account_router.post("/topup")
-async def topup_balance_endpoint(
+
+@wallet_router.post("/topup")
+async def topup_wallet_endpoint(
     cashu_token: str,
     key: ApiKey = Depends(get_key_from_header),
     session: AsyncSession = Depends(get_session),
@@ -39,19 +45,39 @@ async def topup_balance_endpoint(
     return await credit_balance(cashu_token, key, session)
 
 
-@account_router.post("/refund")
-async def refund_balance_endpoint(
+@wallet_router.post("/refund")
+async def refund_wallet_endpoint(
     key: ApiKey = Depends(get_key_from_header),
     session: AsyncSession = Depends(get_session),
 ) -> dict:
-    remaining_balance = key.balance
+    remaining_balance_msats = key.balance
+    
+    if remaining_balance_msats == 0:
+        raise HTTPException(status_code=400, detail="No balance to refund")
+    
+    # Perform refund operation first, before modifying balance
+    if key.refund_address:
+        await refund_balance(remaining_balance_msats, key, session)
+        result = {"recipient": key.refund_address, "msats": remaining_balance_msats}
+    else:
+        # Convert msats to sats for cashu wallet
+        remaining_balance_sats = remaining_balance_msats // 1000
+        if remaining_balance_sats == 0:
+            raise HTTPException(status_code=400, detail="Balance too small to refund (less than 1 sat)")
+        
+        token = await WALLET.send(remaining_balance_sats)
+        result = {"msats": remaining_balance_msats, "recipient": None, "token": token}
+    
+    # Only after successful refund, zero out the balance
     key.balance = 0
     session.add(key)
     await session.commit()
-    if key.refund_address:
-        await refund_balance(remaining_balance, key, session)
-        return {"recipient": key.refund_address, "msats": remaining_balance}
-    else:
-        async with Wallet(nsec=NSEC, mint_urls=[MINT]) as wallet:
-            token = await wallet.send(remaining_balance)
-            return {"msats": remaining_balance, "recipient": None, "token": token}
+    
+    return result
+
+
+@wallet_router.api_route(
+    "/{path:path}", methods=["GET", "POST", "PUT", "DELETE"], include_in_schema=False
+)
+async def wallet_catch_all(path: str):
+    raise HTTPException(status_code=404, detail="Not found check /docs for available endpoints")