message_processor.py

# telegram_summarizer_bot/message_processor.py
"""
Handles aggregation of messages from the database to prepare them for summarization.
"""
import logging
import datetime
from typing import List, Tuple, Optional
import hashlib
from dataclasses import dataclass
import database

logger = logging.getLogger(__name__)

@dataclass
class MessageBatch:
    messages: List[database.FetchedMessage]
    total_tokens: int
    batch_id: str

class MessageProcessor:
    def __init__(self, db_session_factory):
        self.db_session_factory = db_session_factory
        self.max_tokens_per_batch = 4000  # Adjust based on your OpenAI model's context window
        self.estimated_tokens_per_char = 4  # Rough estimate for token calculation

    def _calculate_message_tokens(self, text: str) -> int:
        """Rough estimate of tokens in a text."""
        return len(text) // self.estimated_tokens_per_char

    def _create_message_hash(self, text: str, timestamp: datetime.datetime) -> str:
        """Create a unique hash for a message to detect duplicates."""
        content = f"{text}_{timestamp.isoformat()}"
        return hashlib.md5(content.encode()).hexdigest()

    def _create_batch_id(self, messages: List[database.FetchedMessage]) -> str:
        """Create a unique identifier for a batch of messages."""
        if not messages:
            return ""
        timestamps = [msg.timestamp.isoformat() for msg in messages]
        content = "_".join(timestamps)
        return hashlib.md5(content.encode()).hexdigest()

    async def get_and_aggregate_messages_for_day(
        self, 
        channel_db_id: int, 
        target_date: datetime.date
    ) -> Tuple[Optional[str], List[int]]:
        """
        Fetches messages for a specific channel and date, then aggregates their text.
        Returns a tuple: (aggregated_text, list_of_message_db_ids).
        """
        start_datetime = datetime.datetime.combine(target_date, datetime.time.min, tzinfo=datetime.timezone.utc)
        end_datetime = datetime.datetime.combine(target_date, datetime.time.max, tzinfo=datetime.timezone.utc)
        start_datetime_naive = start_datetime.replace(tzinfo=None)
        end_datetime_naive = end_datetime.replace(tzinfo=None)

        logger.info(f"Aggregating messages for channel DB ID {channel_db_id} for date {target_date.isoformat()}")

        async with self.db_session_factory() as db:
            messages = await database.get_messages_for_summary(db, channel_db_id, start_datetime_naive, end_datetime_naive)

        if not messages:
            logger.info(f"No messages found for channel DB ID {channel_db_id} on {target_date.isoformat()}")
            return None, []

        # Deduplicate messages
        seen_hashes = set()
        unique_messages = []
        for msg in messages:
            if not msg.text:
                continue
            msg_hash = self._create_message_hash(msg.text, msg.timestamp)
            if msg_hash not in seen_hashes:
                seen_hashes.add(msg_hash)
                unique_messages.append(msg)

        # Create batches based on token limits
        batches = []
        current_batch = []
        current_tokens = 0

        for msg in unique_messages:
            msg_tokens = self._calculate_message_tokens(msg.text)
            
            if current_tokens + msg_tokens > self.max_tokens_per_batch and current_batch:
                # Create a new batch
                batch_id = self._create_batch_id(current_batch)
                batches.append(MessageBatch(
                    messages=current_batch.copy(),
                    total_tokens=current_tokens,
                    batch_id=batch_id
                ))
                current_batch = []
                current_tokens = 0
            
            current_batch.append(msg)
            current_tokens += msg_tokens

        # Add the last batch if it has messages
        if current_batch:
            batch_id = self._create_batch_id(current_batch)
            batches.append(MessageBatch(
                messages=current_batch,
                total_tokens=current_tokens,
                batch_id=batch_id
            ))

        # Aggregate text for each batch
        aggregated_texts = []
        message_db_ids = []
        
        for batch in batches:
            batch_texts = []
            for msg in batch.messages:
                batch_texts.append(msg.text)
                message_db_ids.append(msg.id)
            
            # Add batch metadata
            batch_header = f"[Batch {batch.batch_id[:8]} - {len(batch.messages)} messages]"
            aggregated_texts.append(f"{batch_header}\n\n" + "\n\n---\n\n".join(batch_texts))

        full_text_block = "\n\n=== BATCH SEPARATOR ===\n\n".join(aggregated_texts)
        
        logger.info(f"Aggregated {len(message_db_ids)} unique messages into {len(batches)} batches for channel DB ID {channel_db_id}")
        return full_text_block, message_db_ids

    async def get_and_aggregate_messages_for_time_range(
        self,
        channel_ids: List[int],
        start_datetime: datetime.datetime,
        end_datetime: datetime.datetime
    ) -> Tuple[Optional[str], List[int]]:
        """
        Fetches and aggregates messages for multiple channels within a time range.
        Returns a tuple: (aggregated_text, list_of_message_db_ids).
        """
        logger.info(f"Aggregating messages for channels {channel_ids} from {start_datetime} to {end_datetime}")

        async with self.db_session_factory() as db:
            messages = await database.get_messages_for_time_range(db, channel_ids, start_datetime, end_datetime)

        if not messages:
            logger.info(f"No messages found for channels {channel_ids} in the specified time range")
            return None, []

        # Group messages by channel
        messages_by_channel = {}
        for msg in messages:
            if not msg.text:
                continue
            if msg.monitored_channel_id not in messages_by_channel:
                messages_by_channel[msg.monitored_channel_id] = []
            messages_by_channel[msg.monitored_channel_id].append(msg)

        # Process each channel's messages
        all_batches = []
        all_message_ids = []

        for channel_id, channel_messages in messages_by_channel.items():
            # Deduplicate messages for this channel
            seen_hashes = set()
            unique_messages = []
            for msg in channel_messages:
                msg_hash = self._create_message_hash(msg.text, msg.timestamp)
                if msg_hash not in seen_hashes:
                    seen_hashes.add(msg_hash)
                    unique_messages.append(msg)

            # Create batches for this channel
            current_batch = []
            current_tokens = 0

            for msg in unique_messages:
                msg_tokens = self._calculate_message_tokens(msg.text)
                
                if current_tokens + msg_tokens > self.max_tokens_per_batch and current_batch:
                    batch_id = self._create_batch_id(current_batch)
                    all_batches.append(MessageBatch(
                        messages=current_batch.copy(),
                        total_tokens=current_tokens,
                        batch_id=batch_id
                    ))
                    current_batch = []
                    current_tokens = 0
                
                current_batch.append(msg)
                current_tokens += msg_tokens

            if current_batch:
                batch_id = self._create_batch_id(current_batch)
                all_batches.append(MessageBatch(
                    messages=current_batch,
                    total_tokens=current_tokens,
                    batch_id=batch_id
                ))

        # Aggregate text for all batches
        aggregated_texts = []
        
        for batch in all_batches:
            batch_texts = []
            for msg in batch.messages:
                batch_texts.append(msg.text)
                all_message_ids.append(msg.id)
            
            # Add batch metadata
            batch_header = f"[Batch {batch.batch_id[:8]} - {len(batch.messages)} messages]"
            aggregated_texts.append(f"{batch_header}\n\n" + "\n\n---\n\n".join(batch_texts))

        full_text_block = "\n\n=== BATCH SEPARATOR ===\n\n".join(aggregated_texts)
        
        logger.info(f"Aggregated {len(all_message_ids)} unique messages into {len(all_batches)} batches across {len(messages_by_channel)} channels")
        return full_text_block, all_message_ids

    async def mark_messages_processed(self, message_db_ids: list[int]):
        """Marks a list of messages (by their database IDs) as processed for summary."""
        if not message_db_ids:
            return
        async with self.db_session_factory() as db:
            await database.mark_messages_as_processed(db, message_db_ids)
        logger.info(f"Marked {len(message_db_ids)} messages as processed for summary.")


# Example usage:
# async def main_processor_example():
#     # This requires a running DB with data
#     await database.init_db() # Ensure tables exist
#     processor = MessageProcessor(database.AsyncSessionLocal)
     
#     # Assume channel with DB ID 1 exists and has messages for yesterday
#     yesterday = datetime.date.today() - datetime.timedelta(days=1)
#     aggregated_text, msg_ids = await processor.get_and_aggregate_messages_for_day(1, yesterday)
    
#     if aggregated_text:
#         print(f"Aggregated text for channel 1 on {yesterday.isoformat()}:\n{aggregated_text[:500]}...")
#         # await processor.mark_messages_processed(msg_ids) # Example of marking
#     else:
#         print(f"No messages to aggregate for channel 1 on {yesterday.isoformat()}")

# if __name__ == "__main__":
#     import asyncio
#     asyncio.run(main_processor_example())