From 1cd53eb890532d8fc2321f6b02e5cd82a29725a2 Mon Sep 17 00:00:00 2001
From: prosolis <5590409+prosolis@users.noreply.github.com>
Date: Sun, 17 May 2026 12:45:21 -0700
Subject: [PATCH] H5: partial spell-slot refresh on short rest
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Short rest now restores all L1 slots plus floor(level/4) additional slots
at the next-available tier ≥L2, lowest-first. Long rest still does a full
wipe; martials and full-pool casters see no DM change.
---
 internal/plugin/dnd_rest.go      |  25 ++++++
 internal/plugin/dnd_rest_test.go | 133 +++++++++++++++++++++++++++++++
 internal/plugin/dnd_spells.go    |  56 +++++++++++++
 3 files changed, 214 insertions(+)

diff --git a/internal/plugin/dnd_rest.go b/internal/plugin/dnd_rest.go
index 88b5982..14f0c8a 100644
--- a/internal/plugin/dnd_rest.go
+++ b/internal/plugin/dnd_rest.go
@@ -3,6 +3,7 @@ package plugin
 import (
 	"fmt"
 	"math/rand/v2"
+	"sort"
 	"strings"
 	"time"
 )
@@ -105,6 +106,8 @@ func (p *AdventurePlugin) handleDnDShortRest(ctx MessageContext) error {
 	lockoutEnd := now.Add(dndShortRestLockoutHours * time.Hour)
 	c.RestingUntil = &lockoutEnd
 
+	slotsRestored, _ := partialRefreshSpellSlots(ctx.Sender, c.Level)
+
 	if err := SaveDnDCharacter(c); err != nil {
 		return p.SendDM(ctx.Sender, "Couldn't save rest state: "+err.Error())
 	}
@@ -112,12 +115,34 @@ func (p *AdventurePlugin) handleDnDShortRest(ctx MessageContext) error {
 	msg := fmt.Sprintf(
 		"🛌 **Short rest.** You recover **%d HP** (%d→%d / %d).\n_Charges remaining: %d. You're resting — `!zone` and `!expedition` locked for 1 hour._",
 		c.HPCurrent-before, before, c.HPCurrent, c.HPMax, c.ShortRestCharges)
+	if slotLine := dndShortRestSlotLine(slotsRestored); slotLine != "" {
+		msg += "\n_" + slotLine + "_"
+	}
 	if line := dndRestShortFlavorLine(); line != "" {
 		msg += "\n\n_" + line + "_"
 	}
 	return p.SendDM(ctx.Sender, msg)
 }
 
+// dndShortRestSlotLine renders a "Spell slots restored: 2 (L1), 1 (L2)."
+// footer. Returns "" if nothing was restored, suppressing the line for
+// martials and casters already at full.
+func dndShortRestSlotLine(restored map[int]int) string {
+	if len(restored) == 0 {
+		return ""
+	}
+	levels := make([]int, 0, len(restored))
+	for lvl := range restored {
+		levels = append(levels, lvl)
+	}
+	sort.Ints(levels)
+	parts := make([]string, 0, len(levels))
+	for _, lvl := range levels {
+		parts = append(parts, fmt.Sprintf("%d (L%d)", restored[lvl], lvl))
+	}
+	return "Spell slots restored: " + strings.Join(parts, ", ") + "."
+}
+
 // ── Long rest ────────────────────────────────────────────────────────────────
 
 func (p *AdventurePlugin) handleDnDLongRest(ctx MessageContext) error {
diff --git a/internal/plugin/dnd_rest_test.go b/internal/plugin/dnd_rest_test.go
index 263232c..97e8bca 100644
--- a/internal/plugin/dnd_rest_test.go
+++ b/internal/plugin/dnd_rest_test.go
@@ -146,6 +146,139 @@ func TestShortRest_AlreadyFullHP(t *testing.T) {
 	}
 }
 
+// makeRestTestMage builds a wounded mage (so short rest doesn't bail on
+// full-HP) and provisions the class slot pool. The caller is responsible
+// for burning slots before testing the refresh.
+func makeRestTestMage(t *testing.T, uid id.UserID, level int) *DnDCharacter {
+	t.Helper()
+	c := &DnDCharacter{
+		UserID: uid, Race: RaceHuman, Class: ClassMage, Level: level,
+		STR: 8, DEX: 13, CON: 12, INT: 16, WIS: 10, CHA: 12,
+	}
+	conMod := abilityModifier(c.CON)
+	c.HPMax = computeMaxHP(c.Class, conMod, level)
+	c.HPCurrent = 1
+	c.ArmorClass = computeAC(c.Class, abilityModifier(c.DEX))
+	c.ShortRestCharges = level
+	if err := SaveDnDCharacter(c); err != nil {
+		t.Fatal(err)
+	}
+	if err := createAdvCharacter(uid, "mage_rest_test"); err != nil {
+		t.Fatal(err)
+	}
+	if err := setSpellSlotsForLevel(uid, ClassMage, level); err != nil {
+		t.Fatal(err)
+	}
+	return c
+}
+
+func TestPartialRefreshSpellSlots_L5MageToPsAtL2(t *testing.T) {
+	setupRestTestDB(t)
+	uid := id.UserID("@partial_l5:example")
+	makeRestTestMage(t, uid, 5) // L5 mage → 4 L1, 3 L2, 2 L3
+
+	// Burn everything so the refresh has work to do.
+	pool := slotsForClassLevel(ClassMage, 5)
+	for lvl, total := range pool {
+		for i := 0; i < total; i++ {
+			if _, err := consumeSpellSlot(uid, lvl); err != nil {
+				t.Fatal(err)
+			}
+		}
+	}
+
+	restored, err := partialRefreshSpellSlots(uid, 5)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if got := restored[1]; got != pool[1] {
+		t.Errorf("L1 restored = %d, want %d (all)", got, pool[1])
+	}
+	if got := restored[2]; got != 1 {
+		t.Errorf("L2 restored = %d, want 1 (floor(5/4))", got)
+	}
+	if _, ok := restored[3]; ok {
+		t.Errorf("L3 should not have been restored: %v", restored)
+	}
+
+	slots, _ := getSpellSlots(uid)
+	if used := slots[1][1]; used != 0 {
+		t.Errorf("L1 used after refresh = %d, want 0", used)
+	}
+	if used := slots[2][1]; used != pool[2]-1 {
+		t.Errorf("L2 used after refresh = %d, want %d", used, pool[2]-1)
+	}
+	if used := slots[3][1]; used != pool[3] {
+		t.Errorf("L3 used after refresh = %d, want %d (untouched)", used, pool[3])
+	}
+}
+
+func TestPartialRefreshSpellSlots_NonCasterNoop(t *testing.T) {
+	setupRestTestDB(t)
+	uid := id.UserID("@partial_fighter:example")
+	makeRestTestChar(t, uid, 5) // fighter — no slots
+
+	restored, err := partialRefreshSpellSlots(uid, 5)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if restored != nil {
+		t.Errorf("fighter got slots restored: %v", restored)
+	}
+}
+
+func TestPartialRefreshSpellSlots_FullCasterReturnsEmpty(t *testing.T) {
+	setupRestTestDB(t)
+	uid := id.UserID("@partial_full:example")
+	makeRestTestMage(t, uid, 5) // slots already at used=0 from setup
+
+	restored, err := partialRefreshSpellSlots(uid, 5)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if restored != nil {
+		t.Errorf("full mage got slots restored: %v", restored)
+	}
+}
+
+func TestShortRest_RefreshesPartialSlotsForMage(t *testing.T) {
+	setupRestTestDB(t)
+	uid := id.UserID("@short_mage:example")
+	makeRestTestMage(t, uid, 5)
+	pool := slotsForClassLevel(ClassMage, 5)
+	for lvl, total := range pool {
+		for i := 0; i < total; i++ {
+			consumeSpellSlot(uid, lvl)
+		}
+	}
+
+	p := &AdventurePlugin{}
+	if err := p.handleDnDShortRest(MessageContext{Sender: uid}); err != nil {
+		t.Fatal(err)
+	}
+	slots, _ := getSpellSlots(uid)
+	if used := slots[1][1]; used != 0 {
+		t.Errorf("L1 used after short rest = %d, want 0", used)
+	}
+	if used := slots[2][1]; used != pool[2]-1 {
+		t.Errorf("L2 used after short rest = %d, want %d", used, pool[2]-1)
+	}
+}
+
+func TestDndShortRestSlotLine(t *testing.T) {
+	if got := dndShortRestSlotLine(nil); got != "" {
+		t.Errorf("nil → %q, want empty", got)
+	}
+	if got := dndShortRestSlotLine(map[int]int{}); got != "" {
+		t.Errorf("empty → %q, want empty", got)
+	}
+	got := dndShortRestSlotLine(map[int]int{1: 4, 2: 1})
+	want := "Spell slots restored: 4 (L1), 1 (L2)."
+	if got != want {
+		t.Errorf("got %q, want %q", got, want)
+	}
+}
+
 func TestLongRest_RequiresHousingOrInn(t *testing.T) {
 	setupRestTestDB(t)
 	uid := id.UserID("@long_no_house:example")
diff --git a/internal/plugin/dnd_spells.go b/internal/plugin/dnd_spells.go
index 6aa4d74..f2ced06 100644
--- a/internal/plugin/dnd_spells.go
+++ b/internal/plugin/dnd_spells.go
@@ -4,6 +4,7 @@ import (
 	"database/sql"
 	"errors"
 	"fmt"
+	"sort"
 	"strings"
 	"time"
 
@@ -551,6 +552,61 @@ func refreshSpellSlots(userID id.UserID) error {
 	return err
 }
 
+// partialRefreshSpellSlots restores spell slots on short rest. All L1 slots
+// come back, plus floor(charLevel/4) additional slots distributed
+// lowest-first across tiers ≥2. Returns slot_level→count restored so the
+// caller can render a footer.
+func partialRefreshSpellSlots(userID id.UserID, charLevel int) (map[int]int, error) {
+	slots, err := getSpellSlots(userID)
+	if err != nil {
+		return nil, err
+	}
+	restored := map[int]int{}
+	if pair, ok := slots[1]; ok && pair[1] > 0 {
+		if _, err := db.Get().Exec(
+			`UPDATE dnd_spell_slots SET used = 0 WHERE user_id = ? AND slot_level = 1`,
+			string(userID)); err != nil {
+			return nil, err
+		}
+		restored[1] = pair[1]
+	}
+	budget := charLevel / 4
+	if budget <= 0 {
+		if len(restored) == 0 {
+			return nil, nil
+		}
+		return restored, nil
+	}
+	levels := make([]int, 0, len(slots))
+	for lvl, pair := range slots {
+		if lvl >= 2 && pair[1] > 0 {
+			levels = append(levels, lvl)
+		}
+	}
+	sort.Ints(levels)
+	for _, lvl := range levels {
+		if budget <= 0 {
+			break
+		}
+		used := slots[lvl][1]
+		take := used
+		if take > budget {
+			take = budget
+		}
+		if _, err := db.Get().Exec(
+			`UPDATE dnd_spell_slots SET used = MAX(0, used - ?) WHERE user_id = ? AND slot_level = ?`,
+			take, string(userID), lvl); err != nil {
+			return restored, err
+		}
+		restored[lvl] = take
+		budget -= take
+	}
+	if len(restored) == 0 {
+		return nil, nil
+	}
+	return restored, nil
+}
+
 // ── Known spells ─────────────────────────────────────────────────────────────
 
 func addKnownSpell(userID id.UserID, spellID string, source string, prepared bool) error {