chat_session_manager.py

import os # Added import
from dataclasses import dataclass, field
from datetime import datetime
from typing import Optional, List
import uuid # For generating unique session IDs

@dataclass
class ChatMessage:
    """Represents a single message within a chat session."""
    sender: str  # "user", "gemini", "system", "plugin_result", "system_error"
    text: str
    message_id: str = field(default_factory=lambda: str(uuid.uuid4()))
    timestamp: datetime = field(default_factory=datetime.now)
    image_path: Optional[str] = None  # Path to attached/screenshotted image
    ocr_text: Optional[str] = None    # OCR text from the image, if any

    # Future considerations:
    # - Rich content type (e.g., if Gemini sends back images or formatted cards)
    # - For plugin calls: plugin_name, plugin_args, plugin_result_data (if not just text)

    def __str__(self):
        img_info = f" [Image: {os.path.basename(self.image_path)}]" if self.image_path else ""
        ocr_info = f" [OCR Used]" if self.ocr_text else "" # Just indicate OCR was present for this message context
        return f"[{self.timestamp.strftime('%H:%M:%S')}] {self.sender.capitalize()}:{img_info}{ocr_info} {self.text}"

@dataclass
class ChatSession:
    """Represents a single chat conversation."""
    session_id: str = field(default_factory=lambda: str(uuid.uuid4()))
    title: Optional[str] = "New Chat"
    start_time: datetime = field(default_factory=datetime.now)
    last_modified_time: datetime = field(default_factory=datetime.now) # For sorting sessions
    messages: List[ChatMessage] = field(default_factory=list)

    def add_message(self, sender: str, text: str, image_path: Optional[str] = None, ocr_text: Optional[str] = None) -> ChatMessage:
        """Adds a new message to the session, updates last_modified_time, and returns the message."""
        message = ChatMessage(sender=sender, text=text, image_path=image_path, ocr_text=ocr_text)
        self.messages.append(message)
        self.last_modified_time = message.timestamp # Update session's last modified time
        return message

    def get_full_conversation_text(self, include_system=False, max_history=None) -> str:
        """
        Returns a simple string representation of the conversation history,
        suitable for sending to Gemini as context.
        """
        history_str = []
        messages_to_include = self.messages
        if max_history and len(messages_to_include) > max_history:
            messages_to_include = messages_to_include[-max_history:]

        for msg in messages_to_include:
            if msg.sender == "system" and not include_system:
                continue

            # Prefix for image/ocr context for the LLM if it was part of this message turn
            context_prefix = ""
            if msg.image_path: # This message turn included an image
                context_prefix += "[User sent an image"
                if msg.ocr_text:
                    context_prefix += f" with extracted OCR text: \"{msg.ocr_text[:200]}...\"]\n" # Limit OCR text length for context
                else:
                    context_prefix += "]\n"

            history_str.append(f"{msg.sender.capitalize()}: {context_prefix}{msg.text}")
        return "\n".join(history_str)

# --- Database Manager ---
import sqlite3

class DatabaseManager:
    def __init__(self, db_path="gemini_chat_history.sqlite"):
        self.db_path = db_path
        # self.initialize_database() # Call this explicitly after instantiation if needed by MainApp

    def _get_db_connection(self):
        conn = sqlite3.connect(self.db_path)
        conn.row_factory = sqlite3.Row # Access columns by name
        return conn

    def initialize_database(self):
        """Creates database tables if they don't exist."""
        with self._get_db_connection() as conn:
            cursor = conn.cursor()
            # Sessions table
            cursor.execute("""
                CREATE TABLE IF NOT EXISTS sessions (
                    session_id TEXT PRIMARY KEY,
                    title TEXT,
                    start_time TEXT,
                    last_modified_time TEXT
                )
            """)
            # Messages table
            cursor.execute("""
                CREATE TABLE IF NOT EXISTS messages (
                    message_id TEXT PRIMARY KEY,
                    session_id TEXT,
                    sender TEXT,
                    text TEXT,
                    timestamp TEXT,
                    image_path TEXT,
                    ocr_text TEXT,
                    FOREIGN KEY (session_id) REFERENCES sessions (session_id) ON DELETE CASCADE
                )
            """)
            conn.commit()
        print(f"Database initialized/checked at {self.db_path}")

    def save_session(self, session: ChatSession):
        """Saves or updates a chat session in the database."""
        with self._get_db_connection() as conn:
            cursor = conn.cursor()
            cursor.execute("""
                INSERT OR REPLACE INTO sessions (session_id, title, start_time, last_modified_time)
                VALUES (?, ?, ?, ?)
            """, (session.session_id, session.title, session.start_time.isoformat(), session.last_modified_time.isoformat()))
            conn.commit()
            # print(f"Saved session: {session.session_id}") # Optional: for debugging

    def load_all_sessions(self) -> List[ChatSession]:
        """Loads all chat sessions from the database, ordered by last_modified_time."""
        sessions = []
        with self._get_db_connection() as conn:
            cursor = conn.cursor()
            cursor.execute("SELECT * FROM sessions ORDER BY last_modified_time DESC")
            rows = cursor.fetchall()
            for row in rows:
                # Messages are loaded separately when a session is selected
                session = ChatSession(
                    session_id=row["session_id"],
                    title=row["title"],
                    start_time=datetime.fromisoformat(row["start_time"]),
                    last_modified_time=datetime.fromisoformat(row["last_modified_time"]),
                    messages=[] # Messages will be loaded on demand
                )
                sessions.append(session)
        # print(f"Loaded {len(sessions)} sessions from DB.")
        return sessions

    def save_message(self, message: ChatMessage, session_id: str):
        """Saves a single message to the database."""
        with self._get_db_connection() as conn:
            cursor = conn.cursor()
            cursor.execute("""
                INSERT INTO messages (message_id, session_id, sender, text, timestamp, image_path, ocr_text)
                VALUES (?, ?, ?, ?, ?, ?, ?)
            """, (
                message.message_id, session_id, message.sender, message.text,
                message.timestamp.isoformat(), message.image_path, message.ocr_text
            ))
            conn.commit()
            # print(f"Saved message {message.message_id} for session {session_id}") # Optional

    def load_messages_for_session(self, session_id: str) -> List[ChatMessage]:
        """Loads all messages for a given session_id, ordered by timestamp."""
        messages = []
        with self._get_db_connection() as conn:
            cursor = conn.cursor()
            cursor.execute("SELECT * FROM messages WHERE session_id = ? ORDER BY timestamp ASC", (session_id,))
            rows = cursor.fetchall()
            for row in rows:
                message = ChatMessage(
                    sender=row["sender"],
                    text=row["text"],
                    message_id=row["message_id"],
                    timestamp=datetime.fromisoformat(row["timestamp"]),
                    image_path=row["image_path"],
                    ocr_text=row["ocr_text"]
                )
                messages.append(message)
        # print(f"Loaded {len(messages)} messages for session {session_id}")
        return messages

    def delete_session(self, session_id: str):
        """Deletes a session and its associated messages from the database."""
        with self._get_db_connection() as conn:
            cursor = conn.cursor()
            cursor.execute("DELETE FROM messages WHERE session_id = ?", (session_id,))
            cursor.execute("DELETE FROM sessions WHERE session_id = ?", (session_id,))
            conn.commit()
        print(f"Deleted session {session_id} and its messages.")

    def update_session_title(self, session_id: str, new_title: str):
        """Updates the title of a specific session."""
        with self._get_db_connection() as conn:
            cursor = conn.cursor()
            cursor.execute("UPDATE sessions SET title = ?, last_modified_time = ? WHERE session_id = ?",
                           (new_title, datetime.now().isoformat(), session_id))
            conn.commit()
        print(f"Updated title for session {session_id} to '{new_title}'")


# Global instance of DatabaseManager
# This allows MainApp to import and use it directly.
db_manager = DatabaseManager()


if __name__ == '__main__':
    # Test the data classes and DatabaseManager
    # import os # Already imported at the top

    # --- Database Tests ---
    print("\n--- Testing DatabaseManager ---")
    # Use a temporary DB for testing to avoid affecting the main one
    test_db_path = "test_chat_history.sqlite"
    if os.path.exists(test_db_path):
        os.remove(test_db_path) # Clean start for test

    test_db_manager = DatabaseManager(db_path=test_db_path)
    test_db_manager.initialize_database()

    # Create a session
    session1 = ChatSession(title="My First Test Chat")
    test_db_manager.save_session(session1)
    print(f"Session created: ID={session1.session_id}, Title='{session1.title}', Started='{session1.start_time}'")

    # Add some messages
    msg1 = session1.add_message("user", "Hello Gemini!")
    print(f"Added message: {msg1}")

    msg2 = session1.add_message("gemini", "Hello User! How can I help you today?")
    print(f"Added message: {msg2}")

    msg3 = session1.add_message("user", "Can you tell me about large language models?", image_path="/path/to/some/image.png", ocr_text="This is some text from the image.")
    print(f"Added message: {msg3}")

    print("\nFull Session Messages:")
    for msg in session1.messages:
        print(msg)

    print(f"\nSession Title: {session1.title}")
    print(f"Number of messages: {len(session1.messages)}")

    print("\nFormatted Conversation History for LLM:")
    print(session1.get_full_conversation_text())

    print("\nFormatted Conversation History for LLM (max 1 message):")
    print(session1.get_full_conversation_text(max_history=1))


    session2 = ChatSession() # Default title
    print(f"\nSecond session created: ID={session2.session_id}, Title='{session2.title}'")
    session2.add_message("user", "Another topic...")
    print(session2.messages[0])

    # Test __str__ for ChatMessage
    test_msg_no_img = ChatMessage(sender="user", text="Just text.")
    print(f"\nTest Message (no image): {test_msg_no_img}")
    test_msg_with_img = ChatMessage(sender="user", text="Look at this.", image_path="screenshots/my_screen.jpg")
    print(f"Test Message (with image): {test_msg_with_img}")
    test_msg_with_img_ocr = ChatMessage(sender="user", text="And this text?", image_path="screenshots/ocr_img.png", ocr_text="Hello World from OCR")
    print(f"Test Message (with image & OCR): {test_msg_with_img_ocr}")