import time
from functools import partialmethod
from typing import Set

from clarabot_nano.utility.aio_sqlite import AbstractAioSQLiteDatabase


class IndexingStateDatabase(AbstractAioSQLiteDatabase):
    """
    This database stores the structured state of indexing traversal through filesystems from the drive roots.
    The roots here are the resolved paths from the configuration and not the raw values.
    """

    LAST_POLL_ROOT_CONFIG_KEY = "last_poll_root"
    POLL_TRAVERSAL_BEGIN_KEY = "poll_begin"
    POLL_TRAVERSAL_END_KEY = "poll_end"

    def _initialise(self):
        """
        There are two competing goals or durability groups in this database.
            - The "Roots*" table holds the document roots that are in the index database. These must be durable,
              not even a system crash shall loose updates to them after the call returns from this API.
            - All other tables have very frequently changing records of which none is required to be durable.
        To avoid IO waits on volatile transactions, the db will use WAL journaling and NORMAL sync policy. These supply
        the highest throughput with persistence, when the application can choose to ensure durability at any point by
        issuing a checkpoint.
        """
        self._con.execute("PRAGMA journal_mode = WAL;")
        self._con.execute("PRAGMA synchronous = NORMAL;")
        # TODO: research fine-tuning wal_autocheckpoint

    def _setup(self):
        # Key-value tables to store little configuration, metadata or extra state.
        self._con.execute("CREATE TABLE IF NOT EXISTS ConfigKVStore0 (key TEXT PRIMARY KEY, value TEXT NOT NULL)")
        self._con.execute("CREATE TABLE IF NOT EXISTS ConfigKVStore1 (key TEXT PRIMARY KEY, value INTEGER NOT NULL)")
        # This table holds roots that have been used in the index. It is used for tracking the removal of roots
        # and documents from the index and the last poll traversal's time.
        self._con.execute(
            "CREATE TABLE IF NOT EXISTS Roots0 (root TEXT PRIMARY KEY, polled INTEGER NOT NULL DEFAULT 0)"
        )
        # This table holds relative paths under the current root under polling traversal.
        self._con.execute("CREATE TABLE IF NOT EXISTS Paths0 (path TEXT PRIMARY KEY)")
        # This table holds information on files and directories that have gone missing. There is a grace period for
        # removal of these from the index to help managing network drives without constant re-indexing. All records
        # which have an earlier update time than what was stored for the traversal beginning are mass-deleted when
        # a traversal finishes.
        self._con.execute(
            "CREATE TABLE IF NOT EXISTS Losts0 ("
            " root TEXT NOT NULL,"
            " path TEXT NOT NULL,"
            " t INTEGER NOT NULL DEFAULT (strftime('%s','now')),"
            " c INTEGER NOT NULL DEFAULT 0,"
            " u INTEGER NOT NULL,"
            " PRIMARY KEY (root, path)"
            ")"
        )
        self._con.execute("CREATE INDEX IF NOT EXISTS Losts0__root ON Losts0(root)")  # for root removal optimization
        self._con.execute("CREATE INDEX IF NOT EXISTS Losts0__u ON Losts0(u)")  # for removal by recency
        # Directory traversal is the smallest manageable unit of work when synchronizing. (There is no support for
        # incremental or paged directory listing.) The traversal shall record directories to be visited in a following
        # cycle. The indexing (thus the traversal as well) is specifically designed to be done without parallelization.
        # This is the simplest way to ensure the indexer will not hog too much resources ever. The current root of the
        # traversal is saved in the KV table. Once the traversal becomes empty the root has been fully traversed.
        # The fact that an interrupted directory traversal cannot be continued is somewhat alleviated by the fact that
        # newly indexed files will be kept if their update-control data is matching.
        self._con.execute("CREATE TABLE IF NOT EXISTS Traversal0 (path TEXT PRIMARY KEY)")
        # Different change events will populate this and updates will take priority over the sync-polling traversal.
        # The merge field will help coordinate re-processing of a record when it is re-added during processing.
        self._con.execute(
            "CREATE TABLE IF NOT EXISTS PriorityTraversal1 ("
            " root TEXT NOT NULL,"
            " path TEXT NOT NULL,"
            " merge INTEGER NOT NULL DEFAULT 0,"
            " PRIMARY KEY (root, path)"
            ")"
        )
        # self._con.execute("CREATE INDEX IF NOT EXISTS PriorityTraversal0__root ON PriorityTraversal0(root)")
        # Delete old tables if any.
        self._con.execute("DROP TABLE IF EXISTS PriorityTraversal0")

    def _get_roots(self):
        return set(row[0] for row in self._con.execute("SELECT root FROM Roots0").fetchall())

    def _update_roots(self, roots: Set[str]):
        """
        Add new roots the database that may/will get documents indexed under.
        Return the roots that are in the db but not current. (to facilitate their cleanup)
        """
        old_roots = self._get_roots()
        for new_root in roots.difference(old_roots):
            self._con.execute("INSERT INTO Roots0(root) VALUES (?)", (new_root,))
        # return the old roots that are not in the current roots
        old_roots.difference_update(roots)
        return old_roots

    def _delete_root(self, root: str):
        """
        Remove a root from tracking.
        """
        self._con.execute("DELETE FROM Roots0 WHERE root = ?", (root,))
        self._con.execute("DELETE FROM Losts0 WHERE root = ?", (root,))

    def _reset_poll_traversal(self):
        """
        Start a new poll cycle and abandon/clear last_poll_root as well.
        """
        self._start_poll_traversal()
        self._set_config_str(self.LAST_POLL_ROOT_CONFIG_KEY, None)

    def _init_poll_traversal(self):
        if not self._get_config_int(self.POLL_TRAVERSAL_BEGIN_KEY):
            self._set_config_int(self.POLL_TRAVERSAL_BEGIN_KEY, int(time.time()))

    def _get_poll_traversal_roots_to_process(self):
        poll_begin = self._get_config_int(self.POLL_TRAVERSAL_BEGIN_KEY)
        rows = self._con.execute("SELECT root FROM Roots0 WHERE polled != ?", (poll_begin,)).fetchall()
        return set(row[0] for row in rows)

    def _start_poll_traversal_root(self, root):
        self._con.execute("DELETE FROM Paths0")
        self._con.execute("INSERT INTO Paths0(path) VALUES(?)", ("",))
        self._set_config_str(self.LAST_POLL_ROOT_CONFIG_KEY, root)

    def _next_poll_traversal_path(self):
        rows = self._con.execute("SELECT path FROM Paths0 LIMIT 1").fetchall()
        if rows:
            return rows[0][0]

    def _update_poll_traversal_paths(self, path, child_paths):
        if child_paths:
            for child_path in child_paths:
                self._con.execute("INSERT OR IGNORE INTO Paths0(path) VALUES(?)", (child_path,))
        self._con.execute("DELETE FROM Paths0 WHERE path = ?", (path,))

    def _update_lost_record(self, root, path):
        rows = self._con.execute("SELECT t, c FROM Losts0 WHERE root = ? AND path = ?", (root, path)).fetchall()
        if rows:
            t, c = rows[0]
            c += 1
            self._con.execute(
                "UPDATE Losts0 SET c = ?, u = ? WHERE root = ? AND path = ?", (c, int(time.time()), root, path)
            )
            return t, c
        self._con.execute("INSERT INTO Losts0(root, path, u) VALUES(?, ?, ?)", (root, path, t := int(time.time())))
        return t, 0

    def _finish_poll_traversal_root(self, root):
        poll_begin = self._get_config_int(self.POLL_TRAVERSAL_BEGIN_KEY)
        self._con.execute("UPDATE Roots0 SET polled = ? WHERE root = ?", (poll_begin, root))

    def _finish_poll_traversal(self, roots):
        poll_begin = self._get_config_int(self.POLL_TRAVERSAL_BEGIN_KEY)
        for root in roots:
            self._con.execute("DELETE FROM Losts0 WHERE root = ? AND u < ?", (root, poll_begin))
        self._set_config_str(self.LAST_POLL_ROOT_CONFIG_KEY, None)
        poll_end = self._get_config_int(self.POLL_TRAVERSAL_END_KEY)
        if not poll_end or poll_end <= poll_begin:
            self._set_config_int(self.POLL_TRAVERSAL_END_KEY, max(poll_begin + 1, int(time.time())))
            return 0.0  # only return zero when the poll end is actually set
        return max(time.time() - poll_end, 0.1)

    def _start_poll_traversal(self):
        """
        Wait until the current time (poll begin) becomes different from the old one. This is primarily necessary
        for testing where the poll-interval is less than the resolution of this timestamp (one second).
        """
        poll_begin = self._get_config_int(self.POLL_TRAVERSAL_BEGIN_KEY)
        while True:
            self._set_config_int(self.POLL_TRAVERSAL_BEGIN_KEY, t := int(time.time()))
            if poll_begin != t:
                break
            time.sleep(0.25)

    def _next_priority_traversal_item(self):
        rows = self._con.execute("SELECT root, path, merge FROM PriorityTraversal1 LIMIT 1").fetchall()
        if rows:
            return rows[0]

    def _update_priority_traversal_item(self, root, path, merge, child_paths):
        rows = self._con.execute(
            "SELECT merge FROM PriorityTraversal1 WHERE root = ? AND path = ?", (root, path)
        ).fetchall()
        if rows and rows[0][0] == merge:
            # only apply update if the target item has the same merge state (it will be re-processed if not removed)
            if child_paths:
                for child_path in child_paths:
                    self._con.execute(
                        "INSERT OR IGNORE INTO PriorityTraversal1(root, path) VALUES(?, ?)", (root, child_path)
                    )
            self._con.execute("DELETE FROM PriorityTraversal1 WHERE root = ? AND path = ?", (root, path))

    def _add_priority_traversal_items(self, items):
        for root, path in items:
            rows = self._con.execute(
                "SELECT merge FROM PriorityTraversal1 WHERE root = ? AND path = ?", (root, path)
            ).fetchall()
            if rows:
                self._con.execute(
                    "UPDATE PriorityTraversal1 SET merge = ? WHERE root = ? AND path = ?", (rows[0][0] + 1, root, path)
                )
            else:
                self._con.execute("INSERT INTO PriorityTraversal1(root, path) VALUES(?, ?)", (root, path))

    def _get_config_str(self, key):
        rows = self._con.execute("SELECT value FROM ConfigKVStore0 WHERE key = ?", (key,)).fetchall()
        if rows:
            return rows[0][0]

    def _set_config_str(self, key, value):
        if value is None:
            self._con.execute("DELETE FROM ConfigKVStore0 WHERE key = ?", (key,))
        else:
            self._con.execute("REPLACE INTO ConfigKVStore0(key, value) VALUES (?, ?)", (key, value))

    def _get_config_int(self, key):
        rows = self._con.execute("SELECT value FROM ConfigKVStore1 WHERE key = ?", (key,)).fetchall()
        if rows:
            return rows[0][0]

    def _set_config_int(self, key, value):
        if value is None:
            self._con.execute("DELETE FROM ConfigKVStore1 WHERE key = ?", (key,))
        else:
            self._con.execute("REPLACE INTO ConfigKVStore1(key, value) VALUES (?, ?)", (key, value))

    def _reset(self):
        self._con.execute("DELETE FROM ConfigKVStore0")
        self._con.execute("DELETE FROM ConfigKVStore1")
        self._con.execute("DELETE FROM Roots0")
        self._con.execute("DELETE FROM Paths0")
        self._con.execute("DELETE FROM Losts0")
        self._con.execute("DELETE FROM PriorityTraversal1")

    def _checkpoint(self):
        """
        We treat SQLITE_BUSY (1) result as an error, since this process shall be the sole one using the database. It
        can't be busy with anything else.
        """
        if self._con.execute("PRAGMA wal_checkpoint(TRUNCATE);").fetchall()[0][0] != 0:
            raise Exception("WAL checkpoint failed")

    get_roots = partialmethod(AbstractAioSQLiteDatabase._run_job, _get_roots, begin_immediate=False)
    update_roots = partialmethod(AbstractAioSQLiteDatabase._run_job, _update_roots, post_process_fn=_checkpoint)
    delete_root = partialmethod(AbstractAioSQLiteDatabase._run_job, _delete_root, post_process_fn=_checkpoint)
    get_last_poll_traversal_root = partialmethod(
        AbstractAioSQLiteDatabase._run_job, _get_config_str, LAST_POLL_ROOT_CONFIG_KEY, begin_immediate=False
    )
    reset_poll_traversal = partialmethod(AbstractAioSQLiteDatabase._run_job, _reset_poll_traversal)
    init_poll_traversal = partialmethod(AbstractAioSQLiteDatabase._run_job, _init_poll_traversal)
    get_poll_traversal_roots_to_process = partialmethod(
        AbstractAioSQLiteDatabase._run_job, _get_poll_traversal_roots_to_process, begin_immediate=False
    )
    start_poll_traversal_root = partialmethod(AbstractAioSQLiteDatabase._run_job, _start_poll_traversal_root)
    next_poll_traversal_path = partialmethod(
        AbstractAioSQLiteDatabase._run_job, _next_poll_traversal_path, begin_immediate=False
    )
    update_poll_traversal_paths = partialmethod(AbstractAioSQLiteDatabase._run_job, _update_poll_traversal_paths)
    update_lost_record = partialmethod(AbstractAioSQLiteDatabase._run_job, _update_lost_record)
    finish_poll_traversal_root = partialmethod(AbstractAioSQLiteDatabase._run_job, _finish_poll_traversal_root)
    finish_poll_traversal = partialmethod(AbstractAioSQLiteDatabase._run_job, _finish_poll_traversal)
    start_poll_traversal = partialmethod(AbstractAioSQLiteDatabase._run_job, _start_poll_traversal)
    next_priority_traversal_item = partialmethod(
        AbstractAioSQLiteDatabase._run_job, _next_priority_traversal_item, begin_immediate=False
    )
    update_priority_traversal_item = partialmethod(AbstractAioSQLiteDatabase._run_job, _update_priority_traversal_item)
    add_priority_traversal_items = partialmethod(AbstractAioSQLiteDatabase._run_job, _add_priority_traversal_items)
    get_config_str = partialmethod(AbstractAioSQLiteDatabase._run_job, _get_config_str, begin_immediate=False)
    set_config_str = partialmethod(AbstractAioSQLiteDatabase._run_job, _set_config_str)
    get_config_int = partialmethod(AbstractAioSQLiteDatabase._run_job, _get_config_int, begin_immediate=False)
    set_config_int = partialmethod(AbstractAioSQLiteDatabase._run_job, _set_config_int)
    reset = partialmethod(AbstractAioSQLiteDatabase._run_job, _reset, post_process_fn=_checkpoint)
