1 files changed, 68 insertions, 23 deletions

diff --git a/src/app/update.rs b/src/app/update.rs
index 4e21050..4ef1eb1 100644
--- a/src/app/update.rs
+++ b/src/app/update.rs
@@ -1,5 +1,6 @@
 use async_channel::{Receiver, Sender};
 use chrono::NaiveDate;
+use std::collections::HashSet;
 
 use crate::{
     app::{calendar, form, utils, App},
@@ -30,31 +31,75 @@ pub async fn event_handler(rx: Receiver<Msg>, mut app: App) {
                 form::show(&app, event, false).await;
             }
             Msg::AddEvent { new } => {
-                let date = new.date.clone();
-                app.events.push(new);
-                calendar::refresh_date(&app, date);
+                let refresh_dates = add(&mut app, &new);
+                refresh(&app, &refresh_dates)
             }
             Msg::UpdateEvent { old, new } => {
-                let new_date = new.date.clone();
-                match app.events.iter().position(|e| e.id == new.id) {
-                    Some(index) => {
-                        app.events.remove(index);
-                        app.events.push(new);
-                        calendar::refresh_date(&app, new_date);
-                        if old.date != new_date {
-                            calendar::refresh_date(&app, old.date.clone())
-                        }
-                    }
-                    None => println!("Event not found when updating from {:?} to {:?}", old, new),
-                }
-            }
-            Msg::DeleteEvent { event } => match app.events.iter().position(|e| e.id == event.id) {
-                Some(index) => {
-                    app.events.remove(index);
-                    calendar::refresh_date(&app, event.date);
-                }
-                None => println!("Event not found when trying to delete {:?}", event),
-            },
+                let refresh_dates_1 = remove(&mut app, &old);
+                let refresh_dates_2 = add(&mut app, &new);
+                refresh(&app, &refresh_dates_1.union(&refresh_dates_2).map(|d| *d).collect::<HashSet<NaiveDate>>())
+            }
+            Msg::DeleteEvent { event } => {
+                let refresh_dates = remove(&mut app, &event);
+                refresh(&app, &refresh_dates)
+            }
+        }
+    }
+}
+
+/// Remove event and return dates  that should be refreshed.
+fn remove(app: &mut App, event: &Event) -> HashSet<NaiveDate> {
+    if event.repetition.is_some() {
+        match app.repeated_events.iter().position(|e| e.id == event.id) {
+            Some(index) => {
+                app.repeated_events.remove(index);
+                let mut dates = repetition_dates(&app, event);
+                dates.insert(event.date);
+                dates
+            }
+            None => {
+                eprintln!("Event not found when trying to delete {:?}", event);
+                HashSet::new()
+            }
+        }
+    } else {
+        match app.events.iter().position(|e| e.id == event.id) {
+            Some(index) => {
+                app.events.remove(index);
+                HashSet::from([event.date])
+            }
+            None => {
+                eprintln!("Event not found when trying to delete {:?}", event);
+                HashSet::new()
+            }
         }
     }
 }
+
+/// Add event and return dates that should be refreshed.
+fn add(app: &mut App, event: &Event) -> HashSet<NaiveDate> {
+    if event.repetition.is_some() {
+        app.repeated_events.push(event.clone());
+        let mut dates = repetition_dates(&app, event);
+        dates.insert(event.date);
+        dates
+    } else {
+        app.events.push(event.clone());
+        HashSet::from([event.date])
+    }
+}
+
+/// Repetition dates of a repetead event.
+fn repetition_dates(app: &App, event: &Event) -> HashSet<NaiveDate> {
+    let event_reps = event::repetitions_between(&vec!(event.clone()), app.start_date, app.end_date);
+    HashSet::from_iter(event_reps.keys().map(|d| *d))
+}
+
+/// Refresh app for the given dates.
+fn refresh(app: &App, dates: &HashSet<NaiveDate>) {
+    let repetitions = event::repetitions_between(&app.repeated_events, app.start_date, app.end_date);
+
+    for date in dates {
+        calendar::refresh_date(app, *date, &repetitions)
+    }
+}