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}")