feat: Enhance dose tracking functionality in edit window and add punch button support

src/ui_manager.py

@@ -4,7 +4,7 @@ import sys
 import tkinter as tk
 from collections.abc import Callable
 from datetime import datetime
-from tkinter import ttk
+from tkinter import messagebox, ttk
 from typing import Any
 
 from PIL import Image, ImageTk
@@ -124,7 +124,7 @@ class UIManager:
                 variable=symptom_vars[var_name],
             ).grid(row=idx, column=1, sticky="ew")
 
-        # Medicine tracking section with dose buttons
+        # Medicine tracking section (simplified)
         ttk.Label(input_frame, text="Treatment:").grid(
             row=4, column=0, sticky="w", padx=5, pady=2
         )
@@ -132,7 +132,7 @@ class UIManager:
         medicine_frame.grid(row=4, column=1, padx=0, pady=10, sticky="nsew")
         medicine_frame.grid_columnconfigure(0, weight=1)
 
-        # Store medicine variables and dose tracking
+        # Store medicine variables (checkboxes only)
         medicine_vars: dict[str, tuple[tk.IntVar, str]] = {
             "bupropion": (tk.IntVar(value=0), "Bupropion 150/300 mg"),
             "hydroxyzine": (tk.IntVar(value=0), "Hydroxyzine 25mg"),
@@ -140,44 +140,12 @@ class UIManager:
             "propranolol": (tk.IntVar(value=0), "Propranolol 10mg"),
         }
 
-        # Store dose tracking elements for callback assignment later
-        dose_buttons: dict[str, ttk.Button] = {}
-        dose_entries: dict[str, ttk.Entry] = {}
-        dose_displays: dict[str, tk.Text] = {}
-
-        for idx, (med_name, (var, text)) in enumerate(medicine_vars.items()):
-            # Create a sub-frame for each medicine
-            med_sub_frame = ttk.Frame(medicine_frame)
-            med_sub_frame.grid(row=idx, column=0, sticky="ew", padx=5, pady=2)
-            med_sub_frame.grid_columnconfigure(1, weight=1)
-
-            # Checkbox for medicine taken
-            ttk.Checkbutton(med_sub_frame, text=text, variable=var).grid(
-                row=0, column=0, sticky="w"
+        for idx, (_med_name, (var, text)) in enumerate(medicine_vars.items()):
+            # Just checkbox for medicine taken
+            ttk.Checkbutton(medicine_frame, text=text, variable=var).grid(
+                row=idx, column=0, sticky="w", padx=5, pady=2
             )
 
-            # Dose tracking frame
-            dose_frame = ttk.Frame(med_sub_frame)
-            dose_frame.grid(row=0, column=1, sticky="ew", padx=(10, 0))
-            dose_frame.grid_columnconfigure(0, weight=1)
-
-            # Dose entry and button
-            dose_var = tk.StringVar()
-            dose_entry = ttk.Entry(dose_frame, textvariable=dose_var, width=10)
-            dose_entry.grid(row=0, column=0, sticky="ew", padx=(0, 5))
-
-            dose_button = ttk.Button(dose_frame, text=f"Take {med_name.title()}")
-            dose_button.grid(row=0, column=1)
-
-            # Display area for today's doses (read-only)
-            dose_display = tk.Text(dose_frame, height=2, width=30, wrap=tk.WORD)
-            dose_display.grid(row=1, column=0, columnspan=2, sticky="ew", pady=(2, 0))
-            dose_display.config(state=tk.DISABLED)
-
-            dose_buttons[med_name] = dose_button
-            dose_entries[med_name] = dose_entry
-            dose_displays[med_name] = dose_display
-
         # Note and Date fields
         note_var: tk.StringVar = tk.StringVar()
         date_var: tk.StringVar = tk.StringVar()
@@ -204,9 +172,6 @@ class UIManager:
             "frame": main_container,
             "symptom_vars": symptom_vars,
             "medicine_vars": medicine_vars,
-            "dose_buttons": dose_buttons,
-            "dose_entries": dose_entries,
-            "dose_displays": dose_displays,
             "note_var": note_var,
             "date_var": date_var,
         }
@@ -372,9 +337,9 @@ class UIManager:
         )
         vars_dict.update(med_vars)
 
-        # Dose information display (read-only)
+        # Dose information display (editable)
         current_row += 1
-        self._add_dose_display_to_edit(
+        dose_vars = self._add_dose_display_to_edit(
             edit_win,
             current_row,
             {
@@ -384,6 +349,7 @@ class UIManager:
                 "propranolol": prop_doses,
             },
         )
+        vars_dict.update(dose_vars)
 
         # Note field
         current_row += 2  # Account for dose display
@@ -543,11 +509,23 @@ class UIManager:
         button_frame: ttk.Frame = ttk.Frame(parent)
         button_frame.grid(row=row, column=0, columnspan=2, pady=10)
 
-        # Save button
-        ttk.Button(
-            button_frame,
-            text="Save",
-            command=lambda: callbacks["save"](
+        # Save button - create a custom callback to handle dose data
+        def save_with_doses():
+            # Extract dose data from the text widgets
+            dose_data = {}
+            for medicine in ["bupropion", "hydroxyzine", "gabapentin", "propranolol"]:
+                dose_text_key = f"{medicine}_doses_text"
+                if dose_text_key in vars_dict and isinstance(
+                    vars_dict[dose_text_key], tk.Text
+                ):
+                    raw_text = vars_dict[dose_text_key].get(1.0, tk.END).strip()
+                    dose_data[medicine] = self._parse_dose_text(
+                        raw_text, vars_dict["date"].get()
+                    )
+                else:
+                    dose_data[medicine] = ""
+
+            callbacks["save"](
                 parent,
                 vars_dict["date"].get(),
                 vars_dict["depression"].get(),
@@ -559,7 +537,13 @@ class UIManager:
                 vars_dict["gabapentin"].get(),
                 vars_dict["propranolol"].get(),
                 vars_dict["note"].get(),
-            ),
+                dose_data,
+            )
+
+        ttk.Button(
+            button_frame,
+            text="Save",
+            command=save_with_doses,
         ).pack(side="left", padx=5)
 
         # Cancel button
@@ -576,40 +560,169 @@ class UIManager:
 
     def _add_dose_display_to_edit(
         self, parent: tk.Toplevel, row: int, dose_data: dict[str, str]
-    ) -> None:
-        """Add dose information display to edit window."""
-        ttk.Label(parent, text="Recorded Doses:").grid(
+    ) -> dict[str, tk.Text]:
+        """Add comprehensive dose tracking to edit window with punch buttons."""
+        ttk.Label(parent, text="Dose Tracking:").grid(
             row=row, column=0, sticky="w", padx=5, pady=2
         )
 
-        dose_frame = ttk.LabelFrame(parent, text="Today's Doses (Read-Only)")
+        dose_frame = ttk.LabelFrame(parent, text="Medicine Doses")
         dose_frame.grid(row=row, column=1, padx=5, pady=2, sticky="ew")
+        dose_frame.grid_columnconfigure(2, weight=1)
+
+        dose_vars = {}
 
         for idx, (medicine, doses_str) in enumerate(dose_data.items()):
-            ttk.Label(dose_frame, text=f"{medicine.title()}:").grid(
-                row=idx, column=0, sticky="w", padx=5, pady=1
-            )
+            # Medicine label
+            med_label = ttk.Label(dose_frame, text=f"{medicine.title()}:")
+            med_label.grid(row=idx, column=0, sticky="w", padx=5, pady=2)
 
-            # Parse and display doses
+            # Dose entry field for new doses
+            dose_entry_var = tk.StringVar()
+            dose_entry = ttk.Entry(dose_frame, textvariable=dose_entry_var, width=12)
+            dose_entry.grid(row=idx, column=1, sticky="w", padx=5, pady=2)
+
+            # Store entry variable in dose_vars for access from punch button
+            dose_vars[f"{medicine}_entry_var"] = dose_entry_var
+
+            # Display area for existing doses (editable)
+            dose_text = tk.Text(dose_frame, height=3, width=40, wrap=tk.WORD)
+            dose_text.grid(row=idx, column=2, sticky="ew", padx=5, pady=2)
+
+            # Store text widget in dose_vars
+            dose_vars[f"{medicine}_doses_text"] = dose_text
+
+            # Punch button to record dose immediately
+            punch_button = ttk.Button(
+                dose_frame,
+                text=f"Take {medicine.title()}",
+                width=15,
+                command=lambda med=medicine: self._punch_dose_in_edit(med, dose_vars),
+            )
+            punch_button.grid(row=idx, column=3, sticky="w", padx=5, pady=2)
+
+            # Parse and format doses for editing
             if doses_str:
-                doses_display = []
-                for dose_entry in doses_str.split("|"):
-                    if ":" in dose_entry:
-                        timestamp, dose = dose_entry.split(":", 1)
+                formatted_doses = []
+                for dose_entry_str in doses_str.split("|"):
+                    if ":" in dose_entry_str:
+                        timestamp, dose = dose_entry_str.split(":", 1)
                         # Format timestamp for display
                         try:
-                            from datetime import datetime
-
                             dt = datetime.strptime(timestamp, "%Y-%m-%d %H:%M:%S")
                             time_str = dt.strftime("%H:%M")
-                            doses_display.append(f"{time_str}: {dose}")
+                            formatted_doses.append(f"{time_str}: {dose}")
                         except ValueError:
-                            doses_display.append(dose_entry)
+                            formatted_doses.append(dose_entry_str)
 
-                display_text = ", ".join(doses_display) if doses_display else "None"
+                if formatted_doses:
+                    dose_text.insert(1.0, "\n".join(formatted_doses))
+                else:
+                    dose_text.insert(1.0, "No doses recorded")
             else:
-                display_text = "None"
+                dose_text.insert(1.0, "No doses recorded")
 
-            ttk.Label(dose_frame, text=display_text, wraplength=200).grid(
-                row=idx, column=1, sticky="w", padx=5, pady=1
+            # Add help text below the dose display
+            help_label = ttk.Label(
+                dose_frame,
+                text="Format: HH:MM: dose",
+                font=("TkDefaultFont", 8),
+                foreground="gray",
             )
+            help_label.grid(row=idx, column=4, sticky="w", padx=5, pady=2)
+
+        return dose_vars
+
+    def _punch_dose_in_edit(self, medicine_name: str, dose_vars: dict) -> None:
+        """Handle punch dose button in edit window."""
+        dose_entry_var = dose_vars.get(f"{medicine_name}_entry_var")
+        dose_text_widget = dose_vars.get(f"{medicine_name}_doses_text")
+
+        if not dose_entry_var or not dose_text_widget:
+            return
+
+        dose = dose_entry_var.get().strip()
+
+        if not dose:
+            messagebox.showerror(
+                "Error",
+                f"Please enter a dose amount for {medicine_name}",
+            )
+            return
+
+        # Get current time
+        now = datetime.now()
+        time_str = now.strftime("%H:%M")
+
+        # Get current content
+        current_content = dose_text_widget.get(1.0, tk.END).strip()
+
+        # Add new dose entry
+        new_dose_line = f"{time_str}: {dose}"
+
+        if current_content == "No doses recorded" or not current_content:
+            dose_text_widget.delete(1.0, tk.END)
+            dose_text_widget.insert(1.0, new_dose_line)
+        else:
+            dose_text_widget.insert(tk.END, f"\n{new_dose_line}")
+
+        # Clear the entry field
+        dose_entry_var.set("")
+
+        # Show success message
+        messagebox.showinfo(
+            "Success",
+            f"{medicine_name.title()} dose recorded: {dose} at {time_str}",
+        )
+
+    def _parse_dose_text(self, text: str, date: str) -> str:
+        """Parse dose text from edit window back to CSV format."""
+        if not text or text == "No doses recorded":
+            return ""
+
+        lines = text.strip().split("\n")
+        dose_entries = []
+
+        for line in lines:
+            line = line.strip()
+            if ":" in line and line != "No doses recorded":
+                try:
+                    # Try to parse HH:MM: dose format
+                    # Split on ': ' (colon followed by space) to separate time from dose
+                    if ": " in line:
+                        time_part, dose_part = line.split(": ", 1)
+                    else:
+                        # Fallback: split on first colon after HH:MM pattern
+                        colon_indices = [
+                            i for i, char in enumerate(line) if char == ":"
+                        ]
+                        if len(colon_indices) >= 2:
+                            # Take everything up to the second colon as time
+                            second_colon_idx = colon_indices[1]
+                            time_part = line[:second_colon_idx]
+                            dose_part = line[second_colon_idx + 1 :].strip()
+                        else:
+                            continue
+
+                    dose_part = dose_part.strip()
+
+                    # Create timestamp for today
+                    from datetime import datetime
+
+                    time_str = time_part.strip()
+                    # Parse just the time (HH:MM format)
+                    time_obj = datetime.strptime(time_str, "%H:%M")
+
+                    # Create full timestamp with today's date
+                    today = datetime.strptime(date, "%m/%d/%Y")
+                    full_timestamp = today.replace(
+                        hour=time_obj.hour, minute=time_obj.minute, second=0
+                    )
+
+                    timestamp_str = full_timestamp.strftime("%Y-%m-%d %H:%M:%S")
+                    dose_entries.append(f"{timestamp_str}:{dose_part}")
+                except ValueError:
+                    # If parsing fails, skip this line
+                    continue
+
+        return "|".join(dose_entries)