database_manager.py

import sqlite3
import os
from typing import List, Optional
from chat_session_manager import ChatSession, ChatMessage
from datetime import datetime
import uuid # For message_id if ChatMessage doesn't generate it by default (it does now)

DATABASE_NAME = "gemini_agent_history.db"
DATETIME_FORMAT = "%Y-%m-%d %H:%M:%S.%f" # ISO 8601 like format, good for sorting

def get_db_path():
    home = os.path.expanduser("~")
    app_config_dir = os.path.join(home, ".gemini_desktop_agent")
    if not os.path.exists(app_config_dir):
        try: os.makedirs(app_config_dir)
        except OSError: return DATABASE_NAME
    return os.path.join(app_config_dir, DATABASE_NAME)

def get_db_connection():
    db_path = get_db_path()
    conn = sqlite3.connect(db_path, detect_types=sqlite3.PARSE_DECLTYPES | sqlite3.PARSE_COLNAMES)
    conn.row_factory = sqlite3.Row
    return conn

def initialize_database():
    conn = None
    try:
        conn = get_db_connection()
        cursor = conn.cursor()
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS sessions (
                session_id TEXT PRIMARY KEY,
                title TEXT,
                start_time TEXT NOT NULL,
                last_modified_time TEXT NOT NULL
            )
        """)
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS messages (
                message_id TEXT PRIMARY KEY,
                session_id TEXT NOT NULL,
                sender TEXT NOT NULL,
                text TEXT,
                timestamp TEXT NOT NULL,
                image_path TEXT,
                ocr_text TEXT,
                FOREIGN KEY(session_id) REFERENCES sessions(session_id) ON DELETE CASCADE
            )
        """)
        # Add index for faster message loading
        cursor.execute("CREATE INDEX IF NOT EXISTS idx_messages_session_id ON messages (session_id)")
        cursor.execute("CREATE INDEX IF NOT EXISTS idx_messages_timestamp ON messages (timestamp)")
        cursor.execute("CREATE INDEX IF NOT EXISTS idx_sessions_last_modified ON sessions (last_modified_time)")


        conn.commit()
        print(f"Database initialized/verified at {get_db_path()}")
    except sqlite3.Error as e:
        print(f"Database initialization error: {e}")
    finally:
        if conn: conn.close()

def save_session(session: ChatSession):
    conn = None
    try:
        conn = get_db_connection()
        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.strftime(DATETIME_FORMAT),
              session.last_modified_time.strftime(DATETIME_FORMAT)))
        conn.commit()
        print(f"Session '{session.title}' (ID: {session.session_id}) saved/updated.")
    except sqlite3.Error as e:
        print(f"Error saving session {session.session_id}: {e}")
    finally:
        if conn: conn.close()

def save_message(message: ChatMessage, session_id: str):
    conn = None
    try:
        conn = get_db_connection()
        cursor = conn.cursor()
        cursor.execute("""
            INSERT OR REPLACE 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.strftime(DATETIME_FORMAT),
              message.image_path, message.ocr_text))
        conn.commit()
        # Also update the session's last_modified_time
        cursor.execute("""
            UPDATE sessions SET last_modified_time = ? WHERE session_id = ?
        """, (message.timestamp.strftime(DATETIME_FORMAT), session_id))
        conn.commit()
        # print(f"Message {message.message_id} saved for session {session_id}.")
    except sqlite3.Error as e:
        print(f"Error saving message {message.message_id} for session {session_id}: {e}")
    finally:
        if conn: conn.close()

def load_all_sessions() -> List[ChatSession]:
    sessions = []
    conn = None
    try:
        conn = get_db_connection()
        cursor = conn.cursor()
        cursor.execute("SELECT session_id, title, start_time, last_modified_time FROM sessions ORDER BY last_modified_time DESC")
        rows = cursor.fetchall()
        for row in rows:
            # Messages are not loaded here to keep this light; they'll be lazy-loaded or loaded on demand.
            session = ChatSession(
                session_id=row["session_id"],
                title=row["title"],
                start_time=datetime.strptime(row["start_time"], DATETIME_FORMAT),
                last_modified_time=datetime.strptime(row["last_modified_time"], DATETIME_FORMAT),
                messages=[] # Messages will be loaded separately
            )
            sessions.append(session)
        print(f"Loaded {len(sessions)} sessions metadata.")
    except sqlite3.Error as e:
        print(f"Error loading all sessions: {e}")
    finally:
        if conn: conn.close()
    return sessions

def load_messages_for_session(session_id: str) -> List[ChatMessage]:
    messages = []
    conn = None
    try:
        conn = get_db_connection()
        cursor = conn.cursor()
        cursor.execute("""
            SELECT message_id, sender, text, timestamp, image_path, ocr_text
            FROM messages WHERE session_id = ? ORDER BY timestamp ASC
        """, (session_id,))
        rows = cursor.fetchall()
        for row in rows:
            msg = ChatMessage(
                message_id=row["message_id"],
                sender=row["sender"],
                text=row["text"],
                timestamp=datetime.strptime(row["timestamp"], DATETIME_FORMAT),
                image_path=row["image_path"],
                ocr_text=row["ocr_text"]
            )
            messages.append(msg)
        print(f"Loaded {len(messages)} messages for session {session_id}.")
    except sqlite3.Error as e:
        print(f"Error loading messages for session {session_id}: {e}")
    finally:
        if conn: conn.close()
    return messages

def delete_session(session_id: str):
    conn = None
    try:
        conn = get_db_connection()
        cursor = conn.cursor()
        # ON DELETE CASCADE should handle messages, but explicit delete can also be done.
        cursor.execute("DELETE FROM sessions WHERE session_id = ?", (session_id,))
        # cursor.execute("DELETE FROM messages WHERE session_id = ?", (session_id,)) # Redundant due to CASCADE
        conn.commit()
        print(f"Session {session_id} and its messages deleted.")
    except sqlite3.Error as e:
        print(f"Error deleting session {session_id}: {e}")
    finally:
        if conn: conn.close()

def update_session_title(session_id: str, new_title: str):
    conn = None
    try:
        conn = get_db_connection()
        cursor = conn.cursor()
        new_last_modified = datetime.now().strftime(DATETIME_FORMAT)
        cursor.execute("""
            UPDATE sessions SET title = ?, last_modified_time = ?
            WHERE session_id = ?
        """, (new_title, new_last_modified, session_id))
        conn.commit()
        print(f"Session {session_id} title updated to '{new_title}'.")
    except sqlite3.Error as e:
        print(f"Error updating title for session {session_id}: {e}")
    finally:
        if conn: conn.close()

if __name__ == '__main__':
    db_file = get_db_path()
    if os.path.exists(db_file):
        print(f"Deleting existing test database: {db_file}")
        os.remove(db_file)

    print(f"Database will be created/checked at: {db_file}")
    initialize_database()
    print("Database initialization script finished.")

    # Test CRUD
    test_session1 = ChatSession(title="Test Session Alpha")
    save_session(test_session1)
    msg1_1 = test_session1.add_message("user", "Hello from Alpha")
    save_message(msg1_1, test_session1.session_id)
    # Simulate time passing
    import time; time.sleep(0.01)
    msg1_2 = test_session1.add_message("gemini", "Hi Alpha User")
    save_message(msg1_2, test_session1.session_id)

    test_session2 = ChatSession(title="Test Session Beta")
    save_session(test_session2)
    msg2_1 = test_session2.add_message("user", "Hello from Beta", image_path="test.png")
    save_message(msg2_1, test_session2.session_id)

    print("\n--- Loading all sessions ---")
    all_sessions = load_all_sessions()
    for s in all_sessions:
        print(f"Loaded Session: ID={s.session_id}, Title='{s.title}', Start='{s.start_time}', Modified='{s.last_modified_time}'")
        # Load messages for this session
        s.messages = load_messages_for_session(s.session_id)
        for m in s.messages:
            print(f"  Msg: {m.sender} - {m.text[:30]}... (Img: {m.image_path})")

    print("\n--- Updating title for Beta session ---")
    if all_sessions and len(all_sessions) > 1: # Assuming Beta was loaded second, so it's earlier in reversed list
        beta_session_id_to_update = next(s.session_id for s in all_sessions if "Beta" in s.title)
        if beta_session_id_to_update:
            update_session_title(beta_session_id_to_update, "Test Session Beta (Updated)")
            updated_sessions = load_all_sessions()
            for s in updated_sessions:
                 print(f"After Update - Loaded Session: Title='{s.title}', Modified='{s.last_modified_time}'")


    print("\n--- Deleting Alpha session ---")
    if all_sessions:
        alpha_session_id_to_delete = next(s.session_id for s in all_sessions if "Alpha" in s.title)
        if alpha_session_id_to_delete:
            delete_session(alpha_session_id_to_delete)
            remaining_sessions = load_all_sessions()
            print(f"Remaining sessions: {len(remaining_sessions)}")
            for s in remaining_sessions:
                print(f"Remaining: {s.title}")
                # Check if messages were cascaded (should be 0 for the deleted one)
                # msgs = load_messages_for_session(alpha_session_id_to_delete)
                # print(f"Messages for deleted Alpha session: {len(msgs)}") # Should be 0

    print("\n--- Test complete ---")