import sqlite3
import json
import logging
import os
from datetime import datetime

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class MemoryManager:
    def __init__(self, db_path="session_memory.db"):
        self.db_path = db_path
        self._init_db()

    def _init_db(self):
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                cursor.execute("""
                    CREATE TABLE IF NOT EXISTS messages (
                        id INTEGER PRIMARY KEY AUTOINCREMENT,
                        session_id TEXT NOT NULL,
                        role TEXT NOT NULL,
                        content TEXT NOT NULL,
                        timestamp DATETIME DEFAULT CURRENT_TIMESTAMP
                    )
                """)
                cursor.execute("CREATE INDEX IF NOT EXISTS idx_session_id ON messages(session_id)")
                conn.commit()
        except Exception as e:
            logger.error(f"Failed to init DB: {e}")

    def add_message(self, session_id, role, content):
        """Append a message to the session's history."""
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                cursor.execute(
                    "INSERT INTO messages (session_id, role, content) VALUES (?, ?, ?)",
                    (session_id, role, content)
                )
                conn.commit()
        except Exception as e:
            logger.error(f"Failed to add message: {e}")

    def get_history(self, session_id, limit=10):
        """Retrieve the last `limit` messages."""
        try:
            with sqlite3.connect(self.db_path) as conn:
                conn.row_factory = sqlite3.Row
                cursor = conn.cursor()
                # Get last N messages ordered by time ASC
                cursor.execute(
                    """
                    SELECT role, content FROM (
                        SELECT role, content, timestamp 
                        FROM messages 
                        WHERE session_id = ? 
                        ORDER BY id DESC 
                        LIMIT ?
                    ) ORDER BY id ASC
                    """,
                    (session_id, limit)
                )
                rows = cursor.fetchall()
                return [{"role": r["role"], "content": r["content"]} for r in rows]
        except Exception as e:
            logger.error(f"Failed to retrieve history: {e}")
            return []

    def clear_history(self, session_id):
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                cursor.execute("DELETE FROM messages WHERE session_id = ?", (session_id,))
                conn.commit()
        except Exception as e:
            logger.error(f"Failed to clear history: {e}")