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

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

DEFAULT_STATS_DB = Path.home() / ".dracula" / "stats.db"


class Stats:
    """Manages synchronous usage statistics using an SQLite database.

    Tracks the total number of messages, responses, and characters processed
    by the Dracula instance. It ensures data is safely persisted to a local
    database file, preventing the corruption issues common with JSON files.

    Attributes:
        db_path (Path): The file system path to the SQLite statistics database.
    """

    def __init__(self, db_path: str | Path = None):
        """Initializes the Stats instance and ensures the database is ready.

        Args:
            db_path (str | Path, optional): Custom path for the database.
                Defaults to '~/.dracula/stats.db'.
        """
        self.db_path = Path(db_path) if db_path else DEFAULT_STATS_DB
        self.db_path.parent.mkdir(parents=True, exist_ok=True)
        self._init_db()

    def _init_db(self):
        """Creates the stats table and initializes the single tracking row."""
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                cursor.execute(
                    """
                    CREATE TABLE IF NOT EXISTS stats (
                        id INTEGER PRIMARY KEY CHECK (id = 1),
                        total_messages INTEGER DEFAULT 0,
                        total_responses INTEGER DEFAULT 0,
                        total_characters_sent INTEGER DEFAULT 0,
                        total_characters_received INTEGER DEFAULT 0
                    )
                """
                )
                # Ensure the single tracking row exists
                cursor.execute(
                    """
                    INSERT OR IGNORE INTO stats (id, total_messages, total_responses, total_characters_sent, total_characters_received)
                    VALUES (1, 0, 0, 0, 0)
                """
                )
                conn.commit()
        except Exception as e:
            raise ChatException(f"Failed to initialize stats database: {str(e)}")

    def record_message(self, message: str):
        """Updates the statistics for an outgoing user message.

        Args:
            message (str): The literal text sent to the model.
        """
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                cursor.execute(
                    """
                    UPDATE stats 
                    SET total_messages = total_messages + 1,
                        total_characters_sent = total_characters_sent + ?
                    WHERE id = 1
                """,
                    (len(message),),
                )
                conn.commit()
        except Exception as e:
            logger.error(f"Failed to record message stats: {str(e)}")

    def record_response(self, response: str):
        """Updates the statistics for an incoming model response.

        Args:
            response (str): The literal text generated by the model.
        """
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                cursor.execute(
                    """
                    UPDATE stats 
                    SET total_responses = total_responses + 1,
                        total_characters_received = total_characters_received + ?
                    WHERE id = 1
                """,
                    (len(response),),
                )
                conn.commit()
        except Exception as e:
            logger.error(f"Failed to record response stats: {str(e)}")

    def get_stats(self) -> dict:
        """Retrieves the current usage statistics.

        Returns:
            dict: A dictionary containing the tracked integer values.
        """
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                cursor.execute(
                    "SELECT total_messages, total_responses, total_characters_sent, total_characters_received FROM stats WHERE id = 1"
                )
                row = cursor.fetchone()
                return {
                    "total_messages": row[0],
                    "total_responses": row[1],
                    "total_characters_sent": row[2],
                    "total_characters_received": row[3],
                }
        except Exception as e:
            logger.error(f"Failed to retrieve stats: {str(e)}")
            return {
                "total_messages": 0,
                "total_responses": 0,
                "total_characters_sent": 0,
                "total_characters_received": 0,
            }

    def reset(self):
        """Resets all tracked usage statistics to zero."""
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                cursor.execute(
                    """
                    UPDATE stats 
                    SET total_messages = 0, total_responses = 0, 
                        total_characters_sent = 0, total_characters_received = 0 
                    WHERE id = 1
                """
                )
                conn.commit()
            logger.info("Stats have been reset to zero.")
        except Exception as e:
            raise ChatException(f"Failed to reset stats: {str(e)}")


class AsyncStats:
    """Manages asynchronous usage statistics using an aiosqlite database.

    Provides non-blocking tracking of usage metrics, ensuring the event
    loop remains unblocked during database writes.

    Attributes:
        db_path (Path): The file system path to the SQLite statistics database.
    """

    def __init__(self, db_path: str | Path = None):
        """Initializes the AsyncStats instance with the database path.

        Args:
            db_path (str | Path, optional): Custom path for the database.
        """
        self.db_path = Path(db_path) if db_path else DEFAULT_STATS_DB
        self.db_path.parent.mkdir(parents=True, exist_ok=True)

    async def initialize(self):
        """Asynchronously sets up the stats table and initial row."""
        try:
            async with aiosqlite.connect(self.db_path) as db:
                await db.execute(
                    """
                    CREATE TABLE IF NOT EXISTS stats (
                        id INTEGER PRIMARY KEY CHECK (id = 1),
                        total_messages INTEGER DEFAULT 0,
                        total_responses INTEGER DEFAULT 0,
                        total_characters_sent INTEGER DEFAULT 0,
                        total_characters_received INTEGER DEFAULT 0
                    )
                """
                )
                await db.execute(
                    """
                    INSERT OR IGNORE INTO stats (id, total_messages, total_responses, total_characters_sent, total_characters_received)
                    VALUES (1, 0, 0, 0, 0)
                """
                )
                await db.commit()
        except Exception as e:
            raise ChatException(f"Failed to initialize async stats database: {str(e)}")

    async def record_message(self, message: str):
        """Asynchronously updates stats for an outgoing message."""
        try:
            async with aiosqlite.connect(self.db_path) as db:
                await db.execute(
                    """
                    UPDATE stats 
                    SET total_messages = total_messages + 1,
                        total_characters_sent = total_characters_sent + ?
                    WHERE id = 1
                """,
                    (len(message),),
                )
                await db.commit()
        except Exception as e:
            logger.error(f"Failed to record async message stats: {str(e)}")

    async def record_response(self, response: str):
        """Asynchronously updates stats for an incoming response."""
        try:
            async with aiosqlite.connect(self.db_path) as db:
                await db.execute(
                    """
                    UPDATE stats 
                    SET total_responses = total_responses + 1,
                        total_characters_received = total_characters_received + ?
                    WHERE id = 1
                """,
                    (len(response),),
                )
                await db.commit()
        except Exception as e:
            logger.error(f"Failed to record async response stats: {str(e)}")

    async def get_stats(self) -> dict:
        """Asynchronously retrieves current usage statistics."""
        try:
            async with aiosqlite.connect(self.db_path) as db:
                async with db.execute(
                    "SELECT total_messages, total_responses, total_characters_sent, total_characters_received FROM stats WHERE id = 1"
                ) as cursor:
                    row = await cursor.fetchone()
                    return {
                        "total_messages": row[0],
                        "total_responses": row[1],
                        "total_characters_sent": row[2],
                        "total_characters_received": row[3],
                    }
        except Exception as e:
            logger.error(f"Failed to retrieve async stats: {str(e)}")
            return {
                "total_messages": 0,
                "total_responses": 0,
                "total_characters_sent": 0,
                "total_characters_received": 0,
            }

    async def reset(self):
        """Asynchronously resets all tracked usage statistics to zero."""
        try:
            async with aiosqlite.connect(self.db_path) as db:
                await db.execute(
                    """
                    UPDATE stats 
                    SET total_messages = 0, total_responses = 0, 
                        total_characters_sent = 0, total_characters_received = 0 
                    WHERE id = 1
                """
                )
                await db.commit()
            logger.info("Async stats have been reset to zero.")
        except Exception as e:
            raise ChatException(f"Failed to reset async stats: {str(e)}")
