database.py

# telegram_summarizer_bot/database.py
"""
Database setup and asynchronous operations using SQLAlchemy and aiomysql (for SQLite).
Manages channels, messages, and summaries.
"""
import datetime
import logging
from sqlalchemy import create_engine, Column, Integer, String, Text, DateTime, ForeignKey, UniqueConstraint, Boolean, Float
from sqlalchemy.orm import sessionmaker, declarative_base
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession
from sqlalchemy.future import select # For SQLAlchemy 1.4+ style select

import asyncio # For running sync code in async context if needed
import json
import hashlib
from typing import List, Optional

from config import DATABASE_URL

logger = logging.getLogger(__name__)

# Use async engine for compatibility with asyncio applications
# For SQLite, the URL needs to be adjusted for async
if DATABASE_URL.startswith("sqlite:///"):
    ASYNC_DATABASE_URL = DATABASE_URL.replace("sqlite:///", "sqlite+aiosqlite:///")
else:
    ASYNC_DATABASE_URL = DATABASE_URL

# Create async engine
engine = create_async_engine(ASYNC_DATABASE_URL, echo=False) # Set echo=True for SQL logging
AsyncSessionLocal = sessionmaker(
    autocommit=False, autoflush=False, bind=engine, class_=AsyncSession
)

Base = declarative_base()

class MonitoredChannel(Base):
    __tablename__ = "monitored_channels"
    id = Column(Integer, primary_key=True, index=True)
    telegram_chat_id = Column(String, unique=True, index=True, nullable=False) # e.g., @channelname or channel_id
    name = Column(String, nullable=True) # Optional: friendly name or title
    added_by_user_id = Column(Integer, nullable=False)
    date_added = Column(DateTime, default=datetime.datetime.utcnow)

    def __repr__(self):
        return f"<MonitoredChannel(id={self.id}, telegram_chat_id='{self.telegram_chat_id}')>"

class FetchedMessage(Base):
    __tablename__ = "fetched_messages"
    id = Column(Integer, primary_key=True, index=True)
    monitored_channel_id = Column(Integer, ForeignKey("monitored_channels.id"), nullable=False)
    message_id_telegram = Column(Integer, nullable=False) # Telegram's own message ID
    text = Column(Text, nullable=True)
    sender_id = Column(String, nullable=True) # Can be user ID or channel ID if forwarded
    timestamp = Column(DateTime, nullable=False, index=True)
    date_fetched = Column(DateTime, default=datetime.datetime.utcnow)
    processed_for_summary = Column(Boolean, default=False, index=True)
    message_hash = Column(String(32), nullable=True, index=True)  # MD5 hash of message content
    is_duplicate = Column(Boolean, default=False, index=True)  # Flag for duplicate messages

    __table_args__ = (UniqueConstraint('monitored_channel_id', 'message_id_telegram', name='_channel_message_uc'),)

    def __repr__(self):
        return f"<FetchedMessage(id={self.id}, message_id_telegram={self.message_id_telegram}, channel_id={self.monitored_channel_id})>"

class GeneratedSummary(Base):
    __tablename__ = "generated_summaries"
    id = Column(Integer, primary_key=True, index=True)
    monitored_channel_id = Column(Integer, ForeignKey("monitored_channels.id"), nullable=False)
    summary_text = Column(Text, nullable=False)
    date_for_summary = Column(DateTime, nullable=False, index=True) # e.g., the date the messages are for
    date_generated = Column(DateTime, default=datetime.datetime.utcnow)
    openai_model_used = Column(String, nullable=True)
    batch_id = Column(String(32), nullable=True, index=True)  # For tracking which batch of messages was summarized
    token_count = Column(Integer, nullable=True)  # Track token usage
    cost = Column(Float, nullable=True)  # Track API cost
    message_ids = Column(Text, nullable=True)  # JSON array of message IDs included in this summary

    def __repr__(self):
        return f"<GeneratedSummary(id={self.id}, channel_id={self.monitored_channel_id}, date_for_summary={self.date_for_summary})>"

class SummaryCache(Base):
    __tablename__ = "summary_cache"
    id = Column(Integer, primary_key=True, index=True)
    cache_key = Column(String(64), unique=True, index=True)  # Hash of channel_id + date + batch_id
    summary_id = Column(Integer, ForeignKey("generated_summaries.id"), nullable=False)
    date_cached = Column(DateTime, default=datetime.datetime.utcnow)
    expires_at = Column(DateTime, nullable=True)  # Optional cache expiration

    def __repr__(self):
        return f"<SummaryCache(id={self.id}, cache_key={self.cache_key})>"

async def init_db():
    """Initializes the database and creates tables if they don't exist."""
    async with engine.begin() as conn:
        # await conn.run_sync(Base.metadata.drop_all)  # Use with caution: drops all tables
        await conn.run_sync(Base.metadata.create_all)
    print("Database initialized.")

async def get_db_session() -> AsyncSession:
    """Dependency to get a database session."""
    async with AsyncSessionLocal() as session:
        yield session

# --- Channel Operations ---
async def add_channel_to_monitor(db: AsyncSession, channel_identifier: str, user_id: int, channel_name: str = None):
    existing_channel = await db.execute(
        select(MonitoredChannel).where(MonitoredChannel.telegram_chat_id == channel_identifier)
    )
    if existing_channel.scalars().first():
        return None # Already exists
    
    new_channel = MonitoredChannel(
        telegram_chat_id=channel_identifier, 
        name=channel_name,
        added_by_user_id=user_id
    )
    db.add(new_channel)
    await db.commit()
    await db.refresh(new_channel)
    return new_channel

async def remove_channel_from_monitor(db: AsyncSession, channel_identifier: str):
    result = await db.execute(
        select(MonitoredChannel).where(MonitoredChannel.telegram_chat_id == channel_identifier)
    )
    channel_to_delete = result.scalars().first()
    if channel_to_delete:
        # Consider what to do with fetched messages and summaries for this channel.
        # For now, we'll just delete the channel monitoring entry.
        # You might want to cascade deletes or archive data.
        await db.delete(channel_to_delete)
        await db.commit()
        return True
    return False

async def get_monitored_channels(db: AsyncSession):
    result = await db.execute(select(MonitoredChannel))
    return result.scalars().all()

async def get_monitored_channel_by_identifier(db: AsyncSession, channel_identifier: str):
    result = await db.execute(
        select(MonitoredChannel).where(MonitoredChannel.telegram_chat_id == channel_identifier)
    )
    return result.scalars().first()

async def get_monitored_channel_by_id(db: AsyncSession, channel_id: int):
    """Gets a monitored channel by its database ID."""
    result = await db.execute(
        select(MonitoredChannel).where(MonitoredChannel.id == channel_id)
    )
    return result.scalar_one_or_none()

# --- Message Operations ---
async def store_fetched_messages_batch(db: AsyncSession, channel_db_id: int, messages_data: list):
    """
    Stores a batch of messages. messages_data should be a list of dicts.
    Example message_data dict: {'message_id_telegram': ..., 'text': ..., 'sender_id': ..., 'timestamp': ...}
    """
    logger.info(f"Attempting to store batch of {len(messages_data)} messages for channel {channel_db_id}")
    new_messages = []
    for msg_data in messages_data:
        # Check if message already exists to avoid duplicates (optional, if UniqueConstraint is not enough or for logging)
        # This check can be slow for large batches; rely on DB constraint for performance.
        # existing_msg = await db.execute(
        #     select(FetchedMessage).where(
        #         FetchedMessage.monitored_channel_id == channel_db_id,
        #         FetchedMessage.message_id_telegram == msg_data['message_id_telegram']
        #     )
        # )
        # if existing_msg.scalars().first():
        #     continue # Skip if already exists

        logger.debug(f"Creating message object for message_id {msg_data['message_id_telegram']} from channel {channel_db_id}")
        new_messages.append(
            FetchedMessage(
                monitored_channel_id=channel_db_id,
                message_id_telegram=msg_data['message_id_telegram'],
                text=msg_data.get('text'),
                sender_id=str(msg_data.get('sender_id')), # Ensure string
                timestamp=msg_data['timestamp']
            )
        )
    
    if new_messages:
        logger.info(f"Adding {len(new_messages)} new messages to database")
        db.add_all(new_messages)
        try:
            await db.commit()
            logger.info(f"Successfully committed {len(new_messages)} messages to database")
            return len(new_messages) # Return count of newly added messages
        except Exception as e: # Catch integrity errors for duplicates if not pre-checked
            await db.rollback()
            logger.error(f"Error storing messages, possibly duplicates: {e}")
            # Could attempt to insert one by one if batch fails due to a single duplicate
            # For now, we just report the error and rollback.
            return 0 
    else:
        logger.info("No new messages to store")
    return 0


async def get_messages_for_summary(db: AsyncSession, channel_db_id: int, start_date: datetime.datetime, end_date: datetime.datetime):
    result = await db.execute(
        select(FetchedMessage).where(
            FetchedMessage.monitored_channel_id == channel_db_id,
            FetchedMessage.timestamp >= start_date,
            FetchedMessage.timestamp < end_date, # Use < for end_date to cover up to, but not including, the start of the next day
            FetchedMessage.processed_for_summary == False # Or handle this logic elsewhere
        ).order_by(FetchedMessage.timestamp.asc()) # Ensure chronological order for summary
    )
    return result.scalars().all()

async def mark_messages_as_processed(db: AsyncSession, message_ids: list[int]):
    if not message_ids:
        return
    # This is a bit clunky with SQLAlchemy Core update for async.
    # For a simple update, you might need to iterate or use a more direct update statement if your dialect supports it well with async.
    # A more efficient way would be:
    # from sqlalchemy import update
    # stmt = update(FetchedMessage).where(FetchedMessage.id.in_(message_ids)).values(processed_for_summary=True)
    # await db.execute(stmt)
    # However, for simplicity with ORM objects:
    for msg_id in message_ids:
        msg = await db.get(FetchedMessage, msg_id)
        if msg:
            msg.processed_for_summary = True
    await db.commit()


# --- Summary Operations ---
async def store_summary(db: AsyncSession, channel_db_id: int, summary_text: str, date_for_summary: datetime.datetime, model_used: str):
    new_summary = GeneratedSummary(
        monitored_channel_id=channel_db_id,
        summary_text=summary_text,
        date_for_summary=date_for_summary,
        openai_model_used=model_used
    )
    db.add(new_summary)
    await db.commit()
    await db.refresh(new_summary)
    return new_summary

async def get_summary_for_channel_date(db: AsyncSession, channel_db_id: int, summary_date: datetime.date):
    start_of_day = datetime.datetime.combine(summary_date, datetime.time.min)
    end_of_day = datetime.datetime.combine(summary_date, datetime.time.max)
    
    result = await db.execute(
        select(GeneratedSummary).where(
            GeneratedSummary.monitored_channel_id == channel_db_id,
            GeneratedSummary.date_for_summary >= start_of_day,
            GeneratedSummary.date_for_summary <= end_of_day
        ).order_by(GeneratedSummary.date_generated.desc()) # Get the latest if multiple for some reason
    )
    return result.scalars().first()

async def get_messages_for_time_range(
    db: AsyncSession,
    channel_ids: list[int],
    start_datetime: datetime.datetime,
    end_datetime: datetime.datetime
):
    """Fetches messages for multiple channels within a time range."""
    logger.info(f"Querying messages for channels {channel_ids} between {start_datetime} and {end_datetime}")
    
    # Ensure we're using naive UTC datetimes for database comparison
    start_datetime = start_datetime.replace(tzinfo=None)
    end_datetime = end_datetime.replace(tzinfo=None)
    
    result = await db.execute(
        select(FetchedMessage).where(
            FetchedMessage.monitored_channel_id.in_(channel_ids),
            FetchedMessage.timestamp >= start_datetime,
            FetchedMessage.timestamp <= end_datetime,
            FetchedMessage.text.isnot(None)  # Only get messages with text
        ).order_by(FetchedMessage.timestamp.asc())
    )
    messages = result.scalars().all()
    logger.info(f"Found {len(messages)} messages in the specified time range")
    return messages

async def get_latest_message_for_channel(db: AsyncSession, channel_db_id: int) -> FetchedMessage | None:
    """Gets the most recent message stored for a channel."""
    result = await db.execute(
        select(FetchedMessage)
        .where(FetchedMessage.monitored_channel_id == channel_db_id)
        .order_by(FetchedMessage.message_id_telegram.desc())
        .limit(1)
    )
    return result.scalar_one_or_none()

# Add new functions for summary caching and message deduplication
async def store_summary_with_cache(
    db: AsyncSession,
    channel_db_id: int,
    summary_text: str,
    date_for_summary: datetime.datetime,
    model_used: str,
    batch_id: str,
    token_count: int,
    cost: float,
    message_ids: List[int]
) -> GeneratedSummary:
    """Store a summary and its cache entry."""
    # Create the summary
    new_summary = GeneratedSummary(
        monitored_channel_id=channel_db_id,
        summary_text=summary_text,
        date_for_summary=date_for_summary,
        openai_model_used=model_used,
        batch_id=batch_id,
        token_count=token_count,
        cost=cost,
        message_ids=json.dumps(message_ids)
    )
    db.add(new_summary)
    await db.flush()  # Get the ID without committing

    # Create cache entry
    cache_key = hashlib.md5(
        f"{channel_db_id}_{date_for_summary.isoformat()}_{batch_id}".encode()
    ).hexdigest()
    
    cache_entry = SummaryCache(
        cache_key=cache_key,
        summary_id=new_summary.id,
        expires_at=datetime.datetime.utcnow() + datetime.timedelta(days=7)  # Cache for 7 days
    )
    db.add(cache_entry)
    
    await db.commit()
    await db.refresh(new_summary)
    return new_summary

async def get_cached_summary(
    db: AsyncSession,
    channel_db_id: int,
    date_for_summary: datetime.datetime,
    batch_id: str
) -> Optional[GeneratedSummary]:
    """Get a cached summary if it exists and hasn't expired."""
    cache_key = hashlib.md5(
        f"{channel_db_id}_{date_for_summary.isoformat()}_{batch_id}".encode()
    ).hexdigest()
    
    result = await db.execute(
        select(SummaryCache).where(
            SummaryCache.cache_key == cache_key,
            or_(
                SummaryCache.expires_at.is_(None),
                SummaryCache.expires_at > datetime.datetime.utcnow()
            )
        )
    )
    cache_entry = result.scalar_one_or_none()
    
    if cache_entry:
        result = await db.execute(
            select(GeneratedSummary).where(GeneratedSummary.id == cache_entry.summary_id)
        )
        return result.scalar_one_or_none()
    return None

async def mark_messages_as_duplicates(
    db: AsyncSession,
    message_ids: List[int]
) -> None:
    """Mark messages as duplicates."""
    if not message_ids:
        return
    
    await db.execute(
        update(FetchedMessage)
        .where(FetchedMessage.id.in_(message_ids))
        .values(is_duplicate=True)
    )
    await db.commit()

async def get_duplicate_messages(
    db: AsyncSession,
    channel_db_id: int,
    start_date: datetime.datetime,
    end_date: datetime.datetime
) -> List[FetchedMessage]:
    """Get messages that are marked as duplicates."""
    result = await db.execute(
        select(FetchedMessage).where(
            FetchedMessage.monitored_channel_id == channel_db_id,
            FetchedMessage.timestamp >= start_date,
            FetchedMessage.timestamp <= end_date,
            FetchedMessage.is_duplicate == True
        )
    )
    return result.scalars().all()