From a2a2781285f89374765d986e4394fc540f0f53db Mon Sep 17 00:00:00 2001
From: Andreas <andreasruden91@gmail.com>
Date: Sun, 29 Oct 2023 14:03:23 +0100
Subject: [PATCH] Add in storage counter to stash view

---
 d2warehouse/app/main.py             | 32 ++++++++++++++++++++++++++++-
 d2warehouse/app/templates/item.html |  1 +
 2 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/d2warehouse/app/main.py b/d2warehouse/app/main.py
index afd8347..de8c876 100644
--- a/d2warehouse/app/main.py
+++ b/d2warehouse/app/main.py
@@ -34,6 +34,32 @@ def d2_running() -> bool:
     return False
 
 
+def storage_count(item: Item, stash: str) -> int | str:
+    """How many of this item type exist in storage"""
+    db = get_stash_db(stash)
+    if item.is_simple:
+        return db.execute(
+            "SELECT COUNT(id) FROM item WHERE code = ?", (item.code,)
+        ).fetchone()[0]
+    elif item.quality == Quality.UNIQUE:
+        return db.execute(
+            "SELECT COUNT(id) FROM item INNER JOIN item_extra ON id = item_id WHERE code = ? AND unique_id = ?",
+            (item.code, item.unique_id),
+        ).fetchone()[0]
+    elif item.quality == Quality.SET:
+        return db.execute(
+            "SELECT COUNT(id) FROM item INNER JOIN item_extra ON id = item_id WHERE code = ? AND set_id = ?",
+            (item.code, item.set_id),
+        ).fetchone()[0]
+    elif item.is_runeword:
+        return db.execute(
+            "SELECT COUNT(id) FROM item INNER JOIN item_extra ON id = item_id WHERE code = ? AND runeword_id = ?",
+            (item.code, item.runeword_id),
+        ).fetchone()[0]
+    else:
+        return "N/A"
+
+
 def save_path() -> Path:
     if "D2SAVE_PATH" in os.environ:
         path = Path(os.environ["D2SAVE_PATH"])
@@ -76,7 +102,11 @@ def list_stash(stash_name: str):
     stash = parse_stash(stash_data)
 
     return render_template(
-        "list_stash.html", stash_name=stash_name, stash=stash, stash_hash=stash_hash
+        "list_stash.html",
+        stash_name=stash_name,
+        stash=stash,
+        stash_hash=stash_hash,
+        storage_count=lambda x: storage_count(x, stash_name),
     )
 
 
diff --git a/d2warehouse/app/templates/item.html b/d2warehouse/app/templates/item.html
index 3037291..58a3c0b 100644
--- a/d2warehouse/app/templates/item.html
+++ b/d2warehouse/app/templates/item.html
@@ -5,6 +5,7 @@
     {% if item.quality and item.quality >= 5 %}
     <li class="name color-{{item.color}}">{{item.basename}}</li>
     {% endif %}
+    <li>(in storage: {{storage_count(item)}})</li>
     {% if item.stats %}
     {% for stat in item.stats %}
     <li>{{stat}}</li>