import sqlite3
import json
import uuid
import aiosqlite
from pathlib import Path
from .exceptions import ChatException
from .logger import get_logger

logger = get_logger(f"dracula.{__name__}")

DEFAULT_DB_PATH = Path.home() / ".dracula" / "memory.db"


class _MemoryBase:
    """Shared base class for Memory and AsyncMemory.

    Holds common attributes and the database schema. Subclasses implement
    the actual synchronous or asynchronous I/O.
    """

    _SCHEMA = """
        CREATE TABLE IF NOT EXISTS history (
            id         INTEGER PRIMARY KEY AUTOINCREMENT,
            session_id TEXT    NOT NULL,
            role       TEXT    NOT NULL,
            content    TEXT    NOT NULL,
            timestamp  DATETIME DEFAULT CURRENT_TIMESTAMP
        )
    """
    # Applied once when an old database is missing the session_id column.
    _MIGRATION = "ALTER TABLE history ADD COLUMN session_id TEXT NOT NULL DEFAULT 'default'"

    def __init__(
        self,
        db_path: str | Path = None,
        max_messages: int = 10,
        session_id: str = None,
    ):
        self.db_path = Path(db_path) if db_path else DEFAULT_DB_PATH
        self.db_path.parent.mkdir(parents=True, exist_ok=True)
        self.max_messages = max_messages
        self.session_id = session_id or str(uuid.uuid4())


class Memory(_MemoryBase):
    """Manages conversational state and long-term persistence using SQLite.

    The Memory class handles the storage, retrieval, and maintenance of chat
    history between the user and the Gemini model. By utilizing an SQLite
    backend, it ensures data integrity, prevents memory exhaustion (bloat)
    in long-running sessions, and provides structured access to historical
    context.

    A unique session_id scopes every query so that multiple Dracula instances
    can share one database file without their histories bleeding into each other.

    Attributes:
        db_path (Path): The file system path to the SQLite database.
        max_messages (int): The rolling limit for the number of messages
            retained in active context.
        session_id (str): Unique identifier that isolates this instance's
            history from other sessions in the same database.
    """

    def __init__(
        self,
        db_path: str | Path = None,
        max_messages: int = 10,
        session_id: str = None,
    ):
        """Initializes the Memory instance and establishes the database connection.

        Args:
            db_path (str | Path): Filepath for the SQLite database.
                Defaults to ~/.dracula/memory.db.
            max_messages (int): The maximum number of recent messages to
                keep in history. Defaults to 10.
            session_id (str): Optional identifier for this conversation session.
                Allows multiple isolated sessions in one database file.
                Auto-generated if not provided.
        """
        super().__init__(db_path=db_path, max_messages=max_messages, session_id=session_id)
        self._init_db()

    def _init_db(self):
        """Creates the necessary database schema and runs migrations if needed.

        Raises:
            ChatException: If the database connection fails or the schema
                cannot be initialized.
        """
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                cursor.execute(self._SCHEMA)
                # Migrate old databases that lack the session_id column.
                cols = [row[1] for row in cursor.execute("PRAGMA table_info(history)")]
                if "session_id" not in cols:
                    cursor.execute(self._MIGRATION)
                conn.commit()

            logger.debug("SQLite memory system was successfully initialized.")

        except Exception as e:
            raise ChatException(f"Database could not be initialized: {str(e)}")

    def add_message(self, role: str, content: str):
        """Records a new interaction in the conversation history.

        This method persists a message to the database and automatically
        triggers a history trim to maintain the defined message limit.

        Args:
            role (str): The entity sending the message. Typically 'user'
                or 'model'.
            content (str): The literal text content of the message.
        """
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                cursor.execute(
                    "INSERT INTO history (session_id, role, content) VALUES (?, ?, ?)",
                    (self.session_id, role, content),
                )
                conn.commit()
                self._trim_history()

        except Exception as e:
            logger.error(f"An error occurred while saving the message: {e}")

    def _trim_history(self):
        """Enforces the message retention policy by removing oldest records.

        Only affects messages belonging to the current session.
        """
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                cursor.execute(
                    f"""
                    DELETE FROM history
                    WHERE session_id = ? AND id NOT IN (
                        SELECT id FROM history
                        WHERE session_id = ?
                        ORDER BY id DESC LIMIT {self.max_messages}
                    )
                    """,
                    (self.session_id, self.session_id),
                )
                conn.commit()

        except Exception as e:
            logger.error(f"A problem occurred while trimming the memory: {str(e)}")

    def get_history(self) -> list:
        """Retrieves the full chronological conversation history for this session.

        Returns:
            list: A list of dictionaries, where each dictionary contains
                'role' and 'content' keys.
        """
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                cursor.execute(
                    "SELECT role, content FROM history WHERE session_id = ? ORDER BY id ASC",
                    (self.session_id,),
                )
                rows = cursor.fetchall()
                return [{"role": r, "content": c} for r, c in rows]

        except Exception as e:
            logger.error(f"An error occurred while retrieving the history: {str(e)}")
            return []

    def clear(self):
        """Permanently deletes all records for this session.

        Raises:
            ChatException: If the truncation of the history table fails.
        """
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                cursor.execute("DELETE FROM history WHERE session_id = ?", (self.session_id,))
                conn.commit()

            logger.info("The conversation history has been cleared.")

        except Exception as e:
            raise ChatException(f"The memory could not be erased: {str(e)}")

    def compress_to_summary(self, summary: str, keep_last_n: int = 4):
        """Replace old messages with a summary, preserving the most recent messages.

        Args:
            summary (str): The summary text to insert at the start of history.
            keep_last_n (int): Number of most recent messages to keep verbatim.
        """
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                rows = cursor.execute(
                    "SELECT role, content FROM history WHERE session_id=? ORDER BY id ASC",
                    (self.session_id,),
                ).fetchall()

                if len(rows) <= keep_last_n:
                    return

                kept = rows[-keep_last_n:] if keep_last_n > 0 else []
                cursor.execute("DELETE FROM history WHERE session_id=?", (self.session_id,))
                cursor.execute(
                    "INSERT INTO history (session_id, role, content) VALUES (?,?,?)",
                    (self.session_id, "user", f"[Context Summary] {summary}"),
                )
                for role, content in kept:
                    cursor.execute(
                        "INSERT INTO history (session_id, role, content) VALUES (?,?,?)",
                        (self.session_id, role, content),
                    )
                conn.commit()
        except Exception as e:
            logger.error(f"Failed to compress history: {e}")

    def save(self, filepath: str):
        """Exports the current session's history to a JSON-formatted file.

        Args:
            filepath (str): The destination path for the generated JSON file.
        """
        history = self.get_history()
        Path(filepath).write_text(json.dumps(history, indent=4), encoding="utf-8")

    def load(self, filepath: str):
        """Migrates history from a JSON file into the database for this session.

        Clears the current session's history before importing.

        Args:
            filepath (str): The path to the source JSON file.

        Raises:
            ChatException: If the file is inaccessible or the data
                format is invalid.
        """
        try:
            data = json.loads(Path(filepath).read_text(encoding="utf-8"))
            self.clear()
            for msg in data:
                self.add_message(msg["role"], msg["content"])

        except Exception as e:
            raise ChatException(f"Data could not be transferred: {str(e)}")


class AsyncMemory(_MemoryBase):
    """Handles asynchronous conversational state using aiosqlite.

    This class provides non-blocking database operations for chat history
    management, specifically designed for use within AsyncDracula. It
    ensures that the event loop remains unblocked during I/O operations.

    A unique session_id scopes every query so that multiple AsyncDracula
    instances can share one database file without their histories bleeding
    into each other.

    Attributes:
        db_path (Path): The file system path to the SQLite database.
        max_messages (int): The maximum number of messages to retain in
            the rolling history.
        session_id (str): Unique identifier that isolates this instance's
            history from other sessions in the same database.
    """

    def __init__(
        self,
        db_path: str | Path = None,
        max_messages: int = 10,
        session_id: str = None,
    ):
        """Initializes AsyncMemory with database path and message limits.

        Args:
            db_path (str | Path): Filepath for the SQLite database.
                Defaults to ~/.dracula/memory.db.
            max_messages (int): Maximum number of messages to keep in history.
                Defaults to 10.
            session_id (str): Optional identifier for this conversation session.
                Auto-generated if not provided.
        """
        super().__init__(db_path=db_path, max_messages=max_messages, session_id=session_id)

    async def initialize(self):
        """Asynchronously creates the database schema and runs migrations if needed.

        Must be called before any chat operations.

        Raises:
            ChatException: If the database connection or table creation fails.
        """
        try:
            async with aiosqlite.connect(self.db_path) as db:
                await db.execute(self._SCHEMA)
                # Migrate old databases that lack the session_id column.
                async with db.execute("PRAGMA table_info(history)") as cursor:
                    cols = [row[1] for row in await cursor.fetchall()]
                if "session_id" not in cols:
                    await db.execute(self._MIGRATION)
                await db.commit()

            logger.debug("Async SQLite memory system initialized.")

        except Exception as e:
            raise ChatException(f"Failed to initialize async database: {str(e)}")

    async def add_message(self, role: str, content: str):
        """Asynchronously records a new message and trims the history.

        Args:
            role (str): The entity sending the message ('user' or 'model').
            content (str): The literal text content of the message.
        """
        try:
            async with aiosqlite.connect(self.db_path) as db:
                await db.execute(
                    "INSERT INTO history (session_id, role, content) VALUES (?, ?, ?)",
                    (self.session_id, role, content),
                )
                await db.commit()
                await self._trim_history(db)

        except Exception as e:
            logger.error(f"Error adding message to async memory: {str(e)}")

    async def _trim_history(self, db: aiosqlite.Connection):
        """Asynchronously removes oldest records exceeding the max limit for this session.

        Args:
            db (aiosqlite.Connection): The active database connection.
        """
        try:
            await db.execute(
                f"""
                DELETE FROM history
                WHERE session_id = ? AND id NOT IN (
                    SELECT id FROM history
                    WHERE session_id = ?
                    ORDER BY id DESC LIMIT {self.max_messages}
                )
                """,
                (self.session_id, self.session_id),
            )
            await db.commit()

        except Exception as e:
            logger.warning(f"Failed to trim async history: {str(e)}")

    async def get_history(self) -> list:
        """Retrieves the chronological chat history for this session.

        Returns:
            list: A list of dictionaries with 'role' and 'content' keys.
        """
        try:
            async with aiosqlite.connect(self.db_path) as db:
                async with db.execute(
                    "SELECT role, content FROM history WHERE session_id = ? ORDER BY id ASC",
                    (self.session_id,),
                ) as cursor:
                    rows = await cursor.fetchall()
                    return [{"role": r, "content": c} for r, c in rows]
        except Exception as e:
            logger.error(f"Error fetching async history: {str(e)}")
            return []

    async def clear(self):
        """Asynchronously deletes all records for this session.

        Raises:
            ChatException: If the database operation fails.
        """
        try:
            async with aiosqlite.connect(self.db_path) as db:
                await db.execute(
                    "DELETE FROM history WHERE session_id = ?", (self.session_id,)
                )
                await db.commit()
            logger.info("Async conversation history cleared.")
        except Exception as e:
            raise ChatException(f"Failed to clear async memory: {str(e)}")

    async def compress_to_summary(self, summary: str, keep_last_n: int = 4):
        """Asynchronously replace old messages with a summary.

        Args:
            summary (str): The summary text to insert at the start of history.
            keep_last_n (int): Number of most recent messages to keep verbatim.
        """
        try:
            async with aiosqlite.connect(self.db_path) as db:
                async with db.execute(
                    "SELECT role, content FROM history WHERE session_id=? ORDER BY id ASC",
                    (self.session_id,),
                ) as cursor:
                    rows = await cursor.fetchall()

                if len(rows) <= keep_last_n:
                    return

                kept = rows[-keep_last_n:] if keep_last_n > 0 else []
                await db.execute(
                    "DELETE FROM history WHERE session_id=?", (self.session_id,)
                )
                await db.execute(
                    "INSERT INTO history (session_id, role, content) VALUES (?,?,?)",
                    (self.session_id, "user", f"[Context Summary] {summary}"),
                )
                for role, content in kept:
                    await db.execute(
                        "INSERT INTO history (session_id, role, content) VALUES (?,?,?)",
                        (self.session_id, role, content),
                    )
                await db.commit()
        except Exception as e:
            logger.error(f"Failed to compress async history: {e}")

    async def save(self, filepath: str):
        """Exports this session's history to a JSON-formatted file."""
        history = await self.get_history()
        Path(filepath).write_text(json.dumps(history, indent=4), encoding="utf-8")

    async def load(self, filepath: str):
        """Migrates history from a JSON file into the database for this session."""
        try:
            data = json.loads(Path(filepath).read_text(encoding="utf-8"))
            await self.clear()
            for msg in data:
                await self.add_message(msg["role"], msg["content"])
        except Exception as e:
            raise ChatException(f"Data could not be transferred: {str(e)}")