import sqlite3
import traceback
from typing import Any

from velascobot2.utils import ChatData


class _SQLiteContextMgr(object):
    def __init__(
        self,
        database: str = "sqlitedb",
        isolation_level: str = "",
        ignore_exc: bool = False,
    ):
        self.database = database
        self.isolation_level = isolation_level
        self.ignore_exc = ignore_exc
        self.connection = None
        self.cursor = None

    def __enter__(self):
        try:
            self.connection = sqlite3.connect(
                database=self.database, isolation_level=self.isolation_level
            )
            self.cursor = self.connection.cursor()
            return self.cursor
        except Exception as ex:
            traceback.print_exc()
            raise ex

    def __exit__(self, exc_type, exc_val, exc_tb):
        try:
            if exc_type is not None:
                self.connection.rollback()
                return self.ignore_exc
            else:
                self.connection.commit()
        except Exception as ex:
            traceback.print_exc()
            raise ex
        finally:
            self.cursor.close()
            self.connection.close()


class Database:
    def __init__(self, bot_name: str):
        self.db_name = f"{bot_name}.sqlite"
        self._init_database()

    def _init_database(self):
        """creates schema on first startup"""
        with _SQLiteContextMgr(database=self.db_name) as db:
            db.execute("""
                CREATE TABLE IF NOT EXISTS chats (
                    id TEXT PRIMARY KEY,
                    type TEXT NOT NULL,
                    name TEXT NOT NULL,
                    message_period INTEGER DEFAULT 10,
                    answer_chance REAL DEFAULT 0.5,
                    restricted INTEGER DEFAULT 0,
                    silenced INTEGER DEFAULT 0
                )
            """)

            db.execute("""
                CREATE TABLE IF NOT EXISTS vocab (
                    chat_id TEXT NOT NULL,
                    key TEXT NOT NULL,
                    word TEXT NOT NULL,
                    count INTEGER DEFAULT 1,
                    PRIMARY KEY (chat_id, key, word),
                    FOREIGN KEY (chat_id) REFERENCES chats(id)
                )
            """)

            db.execute("""
                CREATE INDEX IF NOT EXISTS vocab_idx
                ON vocab(chat_id, key)
            """)

            db.execute("""
                CREATE TABLE IF NOT EXISTS graphics (
                    chat_id TEXT NOT NULL,
                    type TEXT NOT NULL,
                    graph_id TEXT NOT NULL,
                    count INTEGER DEFAULT 1,
                    PRIMARY KEY (chat_id, graph_id),
                    FOREIGN KEY (chat_id) REFERENCES chats(id)
                )
            """)

            db.execute("""
                CREATE INDEX IF NOT EXISTS graphics_idx
                ON graphics(chat_id)
            """)

    def ensure_chat_exists(
        self, chat_id: str, chat_type: str, chat_name: str
    ) -> bool:
        """
        checks if the chat exists, stores it with defaults if not.
        returns True if newly created, False if already existed
        """
        with _SQLiteContextMgr(database=self.db_name) as db:
            # check existence
            result = db.execute(
                "SELECT id FROM chats WHERE id = ?", (chat_id,)
            ).fetchone()

            if result:
                return False

            # create new chat config entry
            db.execute(
                "INSERT INTO chats (id, type, name) VALUES (?, ?, ?)",
                (chat_id, chat_type, chat_name),
            )

        return True

    def add_words(self, chat_id: str, key: str, word: str, count: int):
        """adds or updates a word ocurrence"""
        with _SQLiteContextMgr(database=self.db_name) as db:
            db.execute(
                "INSERT INTO vocab (chat_id, key, word, count) VALUES (?, ?, ?,"
                " ?) ON CONFLICT(chat_id, key, word) DO UPDATE SET count ="
                " count + ?",
                (chat_id, key, word, count, count),
            )

    def add_graphic(self, chat_id: str, graph_type: str, graph_id: str):
        """adds or updates a graphic ocurrence"""
        with _SQLiteContextMgr(database=self.db_name) as db:
            db.execute(
                "INSERT INTO graphics (chat_id, type, graph_id, count) VALUES"
                " (?, ?, ?, ?) ON CONFLICT(chat_id, graph_id) DO UPDATE"
                " SET count = count + ?",
                (chat_id, graph_type, graph_id, 1, 1),
            )

    def get_next_words(self, chat_id: str, key: str) -> list[tuple[str, int]]:
        """returns list of possible words to follow the key n-gram"""
        with _SQLiteContextMgr(database=self.db_name) as db:
            rows = db.execute(
                "SELECT word, count FROM vocab WHERE chat_id = ? AND key = ?",
                (chat_id, key),
            ).fetchall()

            if not rows:
                return []

            return rows

    def get_chat_vocab_count(self, chat_id: str) -> int:
        with _SQLiteContextMgr(database=self.db_name) as db:
            result = db.execute(
                "SELECT SUM(count) FROM vocab WHERE chat_id = ?", (chat_id,)
            ).fetchone()[0]
        return result or 0

    def get_chat_graph_count(self, chat_id: str) -> int:
        with _SQLiteContextMgr(database=self.db_name) as db:
            result = db.execute(
                "SELECT SUM(count) FROM graphics WHERE chat_id = ?", (chat_id,)
            ).fetchone()[0]
        return result or 0

    def get_graphics(self, chat_id: str):
        with _SQLiteContextMgr(database=self.db_name) as db:
            rows = db.execute(
                "SELECT type, graph_id, count FROM graphics WHERE chat_id = ?",
                (chat_id,),
            ).fetchall()

            if not rows:
                return []

            return rows

    def _update_chat(self, chat_id: str, field: str, value: Any):
        """updates a field in the chats table"""
        with _SQLiteContextMgr(database=self.db_name) as db:
            db.execute(
                f"UPDATE chats SET {field} = ? WHERE id = ?", (value, chat_id)
            )

    def update_chat_name(self, chat_id: str, chat_name: str):
        """updates a field in the chats table"""
        with _SQLiteContextMgr(database=self.db_name) as db:
            db.execute(
                "UPDATE chats SET name = ? WHERE id = ?", (chat_name, chat_id)
            )

    def update_chat_answer_chance(self, chat_id: str, answer_chance: float):
        with _SQLiteContextMgr(database=self.db_name) as db:
            db.execute(
                "UPDATE chats SET answer_chance = ? WHERE id = ?",
                (answer_chance, chat_id),
            )

    def update_chat_message_period(self, chat_id: str, period: int):
        with _SQLiteContextMgr(database=self.db_name) as db:
            db.execute(
                "UPDATE chats SET message_period = ? WHERE id = ?",
                (period, chat_id),
            )

    def update_chat_set_restricted(self, chat_id: str, restricted: bool):
        with _SQLiteContextMgr(database=self.db_name) as db:
            db.execute(
                "UPDATE chats SET restricted = ? WHERE id = ?",
                (int(restricted), chat_id),
            )

    def update_chat_set_silenced(self, chat_id: str, silenced: bool):
        with _SQLiteContextMgr(database=self.db_name) as db:
            db.execute(
                "UPDATE chats SET silenced = ? WHERE id = ?",
                (int(silenced), chat_id),
            )

    def count_graphics(self, chat_id: str) -> int:
        """Returns the number of graphics stored for a given chat id."""
        with _SQLiteContextMgr(database=self.db_name) as db:
            count = (
                db.execute(
                    "SELECT SUM(count) FROM graphics WHERE chat_id = ?",
                    (chat_id,),
                ).fetchone()[0]
                or 0
            )
        return count

    def count_words(self, chat_id: str, key: str) -> int:
        """Returns the number of words that follow a given key for a given chat id."""
        with _SQLiteContextMgr(database=self.db_name) as db:
            count = (
                db.execute(
                    "SELECT SUM(count) FROM vocab WHERE chat_id = ? AND"
                    " key = ?",
                    (chat_id, key),
                ).fetchone()[0]
                or 0
            )
        return count

    def get_all_chats(self) -> list[ChatData]:
        with _SQLiteContextMgr(database=self.db_name) as db:
            rows = db.execute("SELECT * FROM chats").fetchall()
            return [ChatData.from_sql_row(row) for row in rows]

    def get_chat_data(self, chat_id: str) -> ChatData | None:
        with _SQLiteContextMgr(database=self.db_name) as db:
            result = db.execute(
                "SELECT * FROM chats WHERE id = ?", (chat_id,)
            ).fetchone()

        if result:
            return ChatData.from_sql_row(result)
        return None