From 4425f9f50af7aaffef99ad3be8ad37ba2cc87880 Mon Sep 17 00:00:00 2001
From: Andreas <andreasruden91@gmail.com>
Date: Sat, 28 Oct 2023 18:09:22 +0200
Subject: [PATCH] Add takeout from storage

---
 d2warehouse/app/main.py                     | 76 ++++++++++++++++++++-
 d2warehouse/app/templates/list_storage.html | 23 +++++++
 d2warehouse/stash.py                        | 11 +++
 3 files changed, 107 insertions(+), 3 deletions(-)
 create mode 100644 d2warehouse/app/templates/list_storage.html

diff --git a/d2warehouse/app/main.py b/d2warehouse/app/main.py
index bf17885..4ec39a8 100644
--- a/d2warehouse/app/main.py
+++ b/d2warehouse/app/main.py
@@ -1,16 +1,23 @@
 import hashlib
 from flask import Flask, redirect, abort, render_template, request
 from pathlib import Path
+from d2warehouse.item import Item
 from d2warehouse.parser import parse_stash
 import d2warehouse.db as base_db
 from d2warehouse.app.db import get_db, close_db
 import os
 import re
 
+from d2warehouse.stash import StashFullError
+
 STASH_FILES = {
     "softcore": "SharedStashSoftCoreV2.d2i",
     "hardcore": "SharedStashHardCoreV2.d2i",
 }
+DB_FILES = {
+    "softcore": "d2warehouse.softcore.sqlite3",
+    "hardcore": "d2warehouse.hardcore.sqlite3",
+}
 
 
 def save_path() -> Path:
@@ -24,7 +31,15 @@ def save_path() -> Path:
     return path
 
 
-base_db.set_db_path(str(save_path() / "d2warehouse.sqlite3"))
+def get_stash_db(stash):
+    base_db.set_db_path(str(save_path() / DB_FILES[stash]))
+    return get_db()
+
+
+base_db.set_db_path(str(save_path() / DB_FILES["softcore"]))
+base_db.init_db()
+base_db.close_db()
+base_db.set_db_path(str(save_path() / DB_FILES["hardcore"]))
 base_db.init_db()
 base_db.close_db()
 
@@ -53,7 +68,7 @@ def list_stash(stash_name: str):
 
 @app.route("/stash/<stash_name>/store", methods=["POST"])
 def stash_store_items(stash_name: str):
-    if stash_name not in STASH_FILES:
+    if stash_name not in STASH_FILES or stash_name not in DB_FILES:
         abort(404)
     stash_path = save_path() / STASH_FILES[stash_name]
     tmp_path = save_path() / f"{STASH_FILES[stash_name]}.temp"
@@ -85,10 +100,65 @@ def stash_store_items(stash_name: str):
         stash.tabs[tab_idx].remove(item)
     tmp_path.write_bytes(stash.raw())
 
-    db = get_db()
+    db = get_stash_db(stash_name)
     for _, item in items:
         item.write_to_db(db=db)
 
     tmp_path.replace(stash_path)
 
     return redirect(f"/stash/{stash_name}", code=303)
+
+
+@app.route("/storage/<stash_name>")
+def list_storage(stash_name: str):
+    if stash_name not in DB_FILES:
+        abort(404)
+
+    db = get_stash_db(stash_name)
+    items = {}
+    rows = db.execute("SELECT id FROM item WHERE deleted IS NULL").fetchall()
+    for row in rows:
+        items[row["id"]] = Item.load_from_db(row["id"], db)
+
+    return render_template(
+        "list_storage.html", stash_name=stash_name, storage_items=items
+    )
+
+
+@app.route("/storage/<stash_name>/take", methods=["POST"])
+def storage_take_items(stash_name: str):
+    if stash_name not in STASH_FILES or stash_name not in DB_FILES:
+        abort(404)
+
+    stash_path = save_path() / STASH_FILES[stash_name]
+    tmp_path = save_path() / f"{STASH_FILES[stash_name]}.temp"
+    if tmp_path.exists():
+        # TODO: Handle this condition
+        return "temp file exists (BAD)"
+        return 500
+
+    stash_data = stash_path.read_bytes()
+    stash = parse_stash(stash_data)
+
+    # TODO: create backups
+
+    # Write items to temporary stash file
+    db = get_stash_db(stash_name)
+    ids = [y.group(1) for x in request.form.keys() if (y := re.match(r"item_(\d+)", x))]
+    for id in ids:
+        item = Item.load_from_db(id, db)
+        try:
+            stash.add(item)
+        except StashFullError:
+            return "the shared stash does not fit those items", 500
+    tmp_path.write_bytes(stash.raw())
+
+    # Remove items from db
+    for id in ids:
+        db.execute("UPDATE item SET deleted = CURRENT_TIMESTAMP WHERE id = ?", (id,))
+    db.commit()
+
+    # Finalize by replacing real stash file
+    tmp_path.replace(stash_path)
+
+    return redirect(f"/storage/{stash_name}", code=303)
diff --git a/d2warehouse/app/templates/list_storage.html b/d2warehouse/app/templates/list_storage.html
new file mode 100644
index 0000000..7cdb6f6
--- /dev/null
+++ b/d2warehouse/app/templates/list_storage.html
@@ -0,0 +1,23 @@
+<!doctype html>
+<html lang="en">
+  <head>
+    <meta charset="utf-8">
+    <title>Storage</title>
+    <link rel="stylesheet" href="/static/style.css" />
+  <head>
+  <body>
+    <form action="/storage/{{stash_name}}/take" method="POST">
+    <div>
+      <!-- TODO: Include item.html -->
+      {% for db_id, item in storage_items.items() %}
+      <div>
+        <input type="checkbox" name="item_{{db_id}}" value="take" />
+        {{item.name}}
+        ({{db_id}}) 
+      </div>
+      {% endfor %}
+    </div>
+    <input type="submit" value="Take items">
+    </form>
+  </body>
+</html>
diff --git a/d2warehouse/stash.py b/d2warehouse/stash.py
index 8f5bc75..6f97593 100644
--- a/d2warehouse/stash.py
+++ b/d2warehouse/stash.py
@@ -34,6 +34,17 @@ class Stash:
     def raw(self) -> bytes:
         return b"".join(tab.raw() for tab in self.tabs)
 
+    def add(self, item: Item) -> None:
+        for tab in self.tabs:
+            try:
+                tab.add(item)
+                return
+            except StashFullError:
+                pass
+        raise StashFullError(
+            "Could not locate an open spot in the stash to add the item"
+        )
+
 
 class StashTab:
     def __init__(self) -> None: