1 files changed, 83 insertions, 2 deletions

diff --git a/src/queue.rs b/src/queue.rs
index 9e8e1af..109b748 100644
--- a/src/queue.rs
+++ b/src/queue.rs
@@ -124,6 +124,7 @@ pub struct Queue {
     position: usize,
     shuffle: ShuffleMode,
     shuffled_indices: Vec<usize>,
+    original_tracks: Vec<Track>,
 }
 
 impl Default for Queue {
@@ -139,6 +140,7 @@ impl Queue {
             position: 0,
             shuffle: ShuffleMode::None,
             shuffled_indices: Vec::new(),
+            original_tracks: Vec::new(),
         }
     }
 
@@ -162,8 +164,38 @@ impl Queue {
     }
 
     pub fn set_shuffle(&mut self, mode: ShuffleMode) {
-        self.shuffle = mode;
-        self.update_shuffle_indices();
+        match mode {
+            ShuffleMode::None => {
+                self.shuffle = ShuffleMode::None;
+                self.restore_from_original();
+            }
+            ShuffleMode::All => {
+                self.original_tracks = self.tracks.clone();
+                self.shuffle = ShuffleMode::All;
+                self.update_shuffle_indices();
+            }
+            ShuffleMode::Album => {
+                self.original_tracks = self.tracks.clone();
+                self.shuffle = ShuffleMode::Album;
+                self.update_shuffle_indices_album();
+            }
+        }
+    }
+
+    fn restore_from_original(&mut self) {
+        if self.original_tracks.is_empty() {
+            self.shuffled_indices.clear();
+            return;
+        }
+
+        let current_pos = self.position;
+        let played: Vec<Track> = self.tracks[..current_pos].to_vec();
+        let unplayed_original: Vec<Track> = self.original_tracks[current_pos..].to_vec();
+
+        self.tracks = played;
+        self.tracks.extend(unplayed_original);
+        self.position = current_pos;
+        self.shuffled_indices.clear();
     }
 
     fn update_shuffle_indices(&mut self) {
@@ -180,6 +212,55 @@ impl Queue {
         }
     }
 
+    fn update_shuffle_indices_album(&mut self) {
+        let len = self.tracks.len();
+        self.shuffled_indices = (0..len).collect();
+
+        let mut rng = rand::rng();
+        let pos = self.position.min(len.saturating_sub(1));
+
+        if pos == 0 && len > 1 {
+            self.shuffle_album_range(0..len, &mut rng);
+        } else if pos + 1 < len {
+            self.shuffle_album_range(pos + 1..len, &mut rng);
+        }
+    }
+
+    fn shuffle_album_range(&mut self, range: std::ops::Range<usize>, rng: &mut impl rand::Rng) {
+        if range.start >= range.end {
+            return;
+        }
+
+        let mut album_ranges: Vec<(usize, usize)> = Vec::new();
+        let mut current_album: Option<&str> = None;
+        let mut range_start: Option<usize> = None;
+
+        for i in range.clone() {
+            let album = self.tracks.get(i).and_then(|t| t.album.as_deref());
+            match (current_album, range_start) {
+                (None, None) => {
+                    current_album = album;
+                    range_start = Some(i);
+                }
+                (Some(_), Some(start)) if album != current_album => {
+                    album_ranges.push((start, i));
+                    current_album = album;
+                    range_start = Some(i);
+                }
+                _ => {}
+            }
+        }
+        if let Some(start) = range_start {
+            album_ranges.push((start, range.end));
+        }
+
+        for (start, end) in album_ranges {
+            if start + 1 < end {
+                self.shuffled_indices[start..end].shuffle(rng);
+            }
+        }
+    }
+
     pub fn current(&self) -> Option<&Track> {
         if self.tracks.is_empty() {
             None