144 lines
5.1 KiB
Python
144 lines
5.1 KiB
Python
import logging
|
|
from typing import List, Dict, Any, Tuple, Optional
|
|
from datetime import datetime, UTC
|
|
from apis.himalaya import client as himalaya_client
|
|
|
|
class MessageStore:
|
|
"""Store and manage message envelopes"""
|
|
|
|
def __init__(self):
|
|
self.envelopes: List[Dict[str, Any]] = []
|
|
self.metadata_by_id: Dict[int, Dict[str, Any]] = {}
|
|
self.total_messages = 0
|
|
|
|
def load(self, envelopes: List[Dict[str, Any]], sort_ascending: bool = True) -> None:
|
|
"""Load envelopes from Himalaya client and process them"""
|
|
if not envelopes:
|
|
self.envelopes = []
|
|
self.metadata_by_id = {}
|
|
self.total_messages = 0
|
|
return
|
|
|
|
# Sort by date
|
|
envelopes.sort(
|
|
key=lambda x: x.get("date", ""),
|
|
reverse=not sort_ascending,
|
|
)
|
|
|
|
# Group envelopes by month
|
|
grouped_envelopes = []
|
|
months = {}
|
|
|
|
for envelope in envelopes:
|
|
if "id" not in envelope:
|
|
continue
|
|
|
|
# Extract date and determine month group
|
|
date_str = envelope.get("date", "")
|
|
try:
|
|
date = datetime.fromisoformat(date_str.replace("Z", "+00:00"))
|
|
month_key = date.strftime("%B %Y")
|
|
except (ValueError, TypeError):
|
|
month_key = "Unknown Date"
|
|
|
|
# Add month header if this is a new month
|
|
if month_key not in months:
|
|
months[month_key] = True
|
|
grouped_envelopes.append({"type": "header", "label": month_key})
|
|
|
|
# Add the envelope
|
|
grouped_envelopes.append(envelope)
|
|
|
|
# Store metadata for quick access
|
|
envelope_id = int(envelope["id"])
|
|
self.metadata_by_id[envelope_id] = {
|
|
"id": envelope_id,
|
|
"subject": envelope.get("subject", ""),
|
|
"from": envelope.get("from", {}),
|
|
"to": envelope.get("to", {}),
|
|
"cc": envelope.get("cc", {}),
|
|
"date": date_str,
|
|
"index": len(grouped_envelopes) - 1,
|
|
}
|
|
|
|
self.envelopes = grouped_envelopes
|
|
self.total_messages = len(self.metadata_by_id)
|
|
|
|
async def reload(self, sort_ascending: bool = True) -> None:
|
|
"""Reload envelopes from the Himalaya client"""
|
|
envelopes, success = await himalaya_client.list_envelopes()
|
|
if success:
|
|
self.load(envelopes, sort_ascending)
|
|
else:
|
|
logging.error("Failed to reload envelopes")
|
|
|
|
def get_metadata(self, message_id: int) -> Optional[Dict[str, Any]]:
|
|
"""Get metadata for a message by ID"""
|
|
return self.metadata_by_id.get(message_id)
|
|
|
|
def find_next_valid_id(self, current_index: int) -> Tuple[Optional[int], Optional[int]]:
|
|
"""Find the next valid message ID and its index"""
|
|
if not self.envelopes or current_index >= len(self.envelopes) - 1:
|
|
return None, None
|
|
|
|
# Start from current index + 1
|
|
for idx in range(current_index + 1, len(self.envelopes)):
|
|
item = self.envelopes[idx]
|
|
# Skip header items
|
|
if item and item.get("type") != "header" and "id" in item:
|
|
return int(item["id"]), idx
|
|
|
|
return None, None
|
|
|
|
def find_prev_valid_id(self, current_index: int) -> Tuple[Optional[int], Optional[int]]:
|
|
"""Find the previous valid message ID and its index"""
|
|
if not self.envelopes or current_index <= 0:
|
|
return None, None
|
|
|
|
# Start from current index - 1
|
|
for idx in range(current_index - 1, -1, -1):
|
|
item = self.envelopes[idx]
|
|
# Skip header items
|
|
if item and item.get("type") != "header" and "id" in item:
|
|
return int(item["id"]), idx
|
|
|
|
return None, None
|
|
|
|
def get_oldest_id(self) -> int:
|
|
"""Get the ID of the oldest message (first non-header item)"""
|
|
for item in self.envelopes:
|
|
if item and item.get("type") != "header" and "id" in item:
|
|
return int(item["id"])
|
|
return 0
|
|
|
|
def get_newest_id(self) -> int:
|
|
"""Get the ID of the newest message (last non-header item)"""
|
|
for item in reversed(self.envelopes):
|
|
if item and item.get("type") != "header" and "id" in item:
|
|
return int(item["id"])
|
|
return 0
|
|
|
|
def remove_envelope(self, message_id: int) -> None:
|
|
"""Remove an envelope from the store"""
|
|
metadata = self.metadata_by_id.get(message_id)
|
|
if not metadata:
|
|
return
|
|
|
|
index = metadata["index"]
|
|
if 0 <= index < len(self.envelopes):
|
|
# Remove from the envelopes list
|
|
self.envelopes.pop(index)
|
|
|
|
# Remove from metadata dictionary
|
|
del self.metadata_by_id[message_id]
|
|
|
|
# Update indexes for all subsequent messages
|
|
for id_, meta in self.metadata_by_id.items():
|
|
if meta["index"] > index:
|
|
meta["index"] -= 1
|
|
|
|
# Update total message count
|
|
self.total_messages = len(self.metadata_by_id)
|
|
else:
|
|
logging.warning(f"Invalid index {index} for message ID {message_id}")
|