#!/usr/bin/env python3

from collections import namedtuple
from datetime import datetime, timedelta
from hashlib import sha3_256
from importlib.resources import files
import sqlite3
from textwrap import indent
from typing import Optional


with files(__package__).joinpath("schema.sql").open("r") as fp:
    PRELUDE = fp.read()


def requires_conn(f):
    def _helper(self, *args, **kwargs):
        if self._conn is None:
            raise ConnectionError(f"A connection is required for {f.__name__}")
        return f(self, *args, **kwargs)

    return _helper


def fmap(ctor):
    """Fmap a constructor over the records returned from a query function."""

    def _h1(f):
        def _h2(*args, **kwargs):
            res = f(*args, **kwargs)
            if isinstance(res, list):
                return [ctor(*it) for it in res]
            elif isinstance(res, tuple):
                return ctor(*res)

        return _h2

    return _h1


def one(it, *args, **kwargs):
    assert not args
    assert not kwargs
    return it


class StoreError(Exception):
    pass


class LoginError(StoreError):
    pass


class Store(object):
    def __init__(self, path):
        self._path = path
        self._conn: sqlite3.Connection = None

    def _factory(self, cursor, row):
        fields = [column[0] for column in cursor.description]
        cls = namedtuple("Row", fields)
        return cls._make(row)

    def connect(self):
        if not self._conn:
            self._conn = sqlite3.connect(self._path, isolation_level="IMMEDIATE")
            self._conn.row_factory = self._factory
            for hunk in PRELUDE.split("\n\n"):
                try:
                    self._conn.executescript(hunk).fetchall()
                except sqlite3.OperationalError as e:
                    raise RuntimeError(
                        f"Unable to execute startup script:\n{indent(hunk, ' > ')}"
                    ) from e

    @requires_conn
    def commit(self):
        self._conn.commit()

    def close(self):
        if self._conn:
            self.commit()
            self._conn.close()
            self._conn = None

    ################################################################################
    # Users

    @requires_conn
    def try_create_user(self, username, email, password, group_id=10, status_id=-2):
        """Attempt to create a new user.

        :param username: The name of the user to be created.
        :param email: The email of the user to be created.
        :param password: The (hopefully salted!) plain text password for the user. Will be hashed before storage.
        :param group_id: The numeric ID of a group to assign the user to. Default 10 AKA normal user.
        :param status_id: The numeric ID of the status to assign the user to. Default -2 AKA email verification required.

        """

        digest = sha3_256()
        digest.update(password.encode("utf-8"))
        digest = digest.hexdigest()
        return self._conn.execute(
            "INSERT INTO users (name, email, hash, group_id, status_id) VALUES (?, ?, ?, ?, ?) RETURNING id, status_id",
            [username, email, digest, group_id, status_id],
        ).fetchone()

    @requires_conn
    def fetch_user(self, uid: int):
        return self._conn.execute("SELECT * FROM users WHERE id = ?", [uid]).fetchone()

    @fmap(one)
    @requires_conn
    def fetch_user_priority(self, uid: int):
        return self._conn.execute(
            "SELECT priority FROM groups g INNER JOIN users u ON g.id = u.group_id WHERE u.id = ?",
            [uid],
        ).fetchone()

    @requires_conn
    def list_users(self):
        return self._conn.execute("SELECT id, name FROM users").fetchall()

    @fmap(one)
    @requires_conn
    def fetch_user_status(self, user_status_id: int):
        """Fetch a user status by ID"""

        return self._conn.execute(
            "SELECT id, name FROM user_statuses WHERE id = ?", [user_status_id]
        ).fetchone()

    ################################################################################
    # Sessions / 'keys'

    @fmap(one)
    @requires_conn
    def _create_session(
        self, uid: int, ttl: Optional[timedelta] = None, name: Optional[str] = None
    ):
        return self._conn.execute(
            "INSERT INTO user_keys (user_id, name, expiration) VALUES (?1, ?2, ?3) RETURNING (id)",
            [uid, name, (datetime.now() + ttl).isoformat() if ttl else None],
        ).fetchone()

    @requires_conn
    def try_login(self, username: str, password: str, ttl: timedelta) -> Optional[str]:
        """Given a username and an (unsecured) password, attempt to authenticate the named user.

        If successful, return the ID of a new session/key for that user.

        """

        digest = sha3_256()
        digest.update(password.encode("utf-8"))
        digest = digest.hexdigest()
        res = self._conn.execute(
            "SELECT id, status_id FROM users WHERE (name=?1 AND hash=?2) OR (email=?1 AND hash=?2) LIMIT 1",
            [username, digest],
        ).fetchone()
        if not res:
            return

        uid, status = res
        if status > 0:
            return self._create_session(uid, ttl, "web session")

        else:
            _, status = self.fetch_user_status(status)
            raise LoginError(status)

    @requires_conn
    def create_key(self, kid: str, ttl, name: Optional[str] = None) -> Optional[str]:
        """Given an _existing_ login session, create a new key.

        This allows the user to create more or less permanent API keys associated with their identity.

        """

        if uid := self.try_key(kid):
            return self._create_session(uid, ttl, name)

    @requires_conn
    def list_keys(self, uid: int):
        return [(id, name, datetime.fromisoformat(exp) if exp else None)
         for id, name, exp in self._conn.execute(
            "SELECT id, name, expiration FROM user_keys WHERE user_id = ?1 AND name != 'web session'", [uid]
        ).fetchall()]

    @requires_conn
    def fetch_key(self, kid) -> tuple:
        return self._conn.execute(
            "SELECT * FROM user_keys WHERE id = ?", [kid]
        ).fetchone()

    @fmap(one)
    @requires_conn
    def try_key(self, kid: str):
        """Try to find the mapped user for a session."""

        return self._conn.execute(
            "SELECT user_id FROM user_keys WHERE expiration IS NULL OR unixepoch(expiration) > unixepoch('now') and id = ?",
            [kid],
        ).fetchone()

    @requires_conn
    def refresh_key(self, kid: str, ttl: timedelta):
        """Automagically renew an API key which is still in use.

        Mostly intended for dealing with web sessions which should implicitly extend, but which use the same mechanism as API keys.

        """

        self._conn.execute(
            "UPDATE user_keys SET expiration = ? WHERE id = ?",
            [(datetime.now() + ttl).isoformat(), kid],
        )

    @requires_conn
    def delete_key(self, uid: int, kid: str):
        """Remove a session/key; equivalent to logout."""

        self._conn.execute(
            "DELETE FROM user_keys WHERE user_id = ?1 and id = ?2", [uid, kid]
        )

    ################################################################################
    # Printers
    #
    # Printers represent connections to OctoPrint instances controlling physical machines.

    @fmap(one)
    @requires_conn
    def try_create_printer(self, name: str, url: str, api_key: str):
        self._conn.execute(
            "INSERT INTO printers (name, url, api_key, status_id) VALUES (?1, ?2, ?3, 0) RETURNING id",
            [name, url, api_key],
        ).fetchone()

    @requires_conn
    def fetch_printer(self, printer_id: int):
        return self._conn.execute(
            """
            SELECT
               p.id
             , p.name
             , p.url
             , p.api_key
             , p.last_poll_date
             , s.name as status
            FROM printers p
            INNER JOIN printer_statuses s
            ON p.status_id = s.id
            WHERE p.id = ?1
            """,
            [printer_id],
        ).fetchone()

    @requires_conn
    def list_printers(self):
        return self._conn.execute(
            "SELECT p.id, p.name, p.url, p.api_key, p.last_poll_date, s.name as status FROM printers p INNER JOIN printer_statuses s ON p.status_id = s.id"
        ).fetchall()

    @fmap(one)
    @requires_conn
    def list_idle_printers(self):
        return self._conn.execute(
            """
            SELECT p.id
            FROM printers p
            LEFT JOIN (SELECT id, printer_id FROM jobs WHERE finished_at IS NULL) j ON p.id = j.printer_id
            INNER JOIN printer_statuses s ON p.status_id = s.id
            WHERE j.id IS NULL AND s.name = 'idle'
            """
        ).fetchall()

    @requires_conn
    def update_printer_status(self, printer_id, state: str):
        (status_id,) = self._conn.execute(
            "SELECT id FROM printer_statuses WHERE name = ?1", [state]
        ).fetchone()

        self._conn.execute(
            "UPDATE printers SET status_id = ?2, last_poll_date = datetime('now') WHERE id = ?1",
            [printer_id, status_id],
        )

    ################################################################################
    # Files
    #
    # A record of local files on disk, and the users who own then.
    @fmap(one)
    @requires_conn
    def create_file(self, uid: int, name: str, path: str) -> int:
        return self._conn.execute(
            "INSERT INTO files (user_id, filename, path, upload_date) VALUES (?, ?, ?, datetime('now')) RETURNING (id)",
            [uid, name, path],
        ).fetchone()

    @requires_conn
    def list_files(self, uid: int):
        return self._conn.execute(
            "SELECT * FROM files WHERE user_id = ?", [uid]
        ).fetchall()

    @requires_conn
    def delete_file(self, uid: int, fid: int):
        self._conn.execute("DELETE FROM files WHERE user_id = ? AND id = ?", [uid, fid])

    @requires_conn
    def fetch_file(self, uid: int, fid: int):
        return self._conn.execute("SELECT * FROM files WHERE user_id = ?1 AND id = ?2", [uid, fid]).fetchone()

    ################################################################################
    # Job
    #
    # A request by a user for a given file to be printed.
    @fmap(one)
    @requires_conn
    def create_job(self, uid: int, fid: int, relative_priority: int = 0):
        """Create a job mapped to a user with a file to print and a priority.

        Note that the user may provide a sub-priority within their group queue. This allows users to create jobs with
        higher priority than existing jobs as a means of controlling the queue order.

        May want a different (eg. more explicit) queue ordering mechanism here. Emulating the Netflix queue behavior of
        being able to drag some jobs ahead of others? How to model that?

        """
        assert 0 <= relative_priority <= 9
        return self._conn.execute(
            "INSERT INTO jobs (user_id, file_id, priority) VALUES (?, ?, ?) RETURNING (id)",
            [uid, fid, self.fetch_user_priority(uid) + relative_priority],
        ).fetchone()

    @requires_conn
    def list_jobs(self, uid: Optional[int] = None):
        """Enumerate jobs in priority order."""
        cond = f"user_id = {uid}" if uid else "TRUE"
        return self._conn.execute(
            f"""
            SELECT * FROM jobs
            WHERE {cond}
            """,
            [],
        ).fetchall()

    @requires_conn
    def list_jobs_by_file(self, fid: int):
        return self._conn.execute(
            f"""
            SELECT * FROM jobs
            WHERE file_id = ?1
            """,
            [fid],
        ).fetchall()

    @requires_conn
    def list_job_queue(self, uid: Optional[int] = None):
        """Enumerate jobs in priority order. Note: ignores completed jobs."""
        cond = f"user_id = {uid}" if uid else "TRUE"
        return self._conn.execute(
            f"""
            SELECT * FROM jobs
            WHERE finished_at IS NULL AND {cond}
            ORDER BY priority DESC
            """,
            [],
        ).fetchall()

    @requires_conn
    def list_job_history(self, uid: Optional[int] = None):
        """Enumerate jobs in priority order. Note: ignores completed jobs."""
        cond = f"user_id = {uid}" if uid else "TRUE"
        return self._conn.execute(
            f"""
            SELECT * FROM jobs
            WHERE finished_at IS NOT NULL AND {cond}
            ORDER BY datetime(finished_at) DESC
            """,
            [],
        ).fetchall()

    @requires_conn
    def list_mapped_jobs(self):
        """Scheduler detail. List mapped but not started jobs."""

        return self._conn.execute(
            """
            SELECT * FROM jobs
            WHERE started_at IS NULL AND printer_id IS NOT NULL
            """,
            [],
        ).fetchall()

    @requires_conn
    def list_running_jobs(self):
        """Scheduler detail. List running jobs.

        Note that jobs for which cancellation has been requested but which HAVE
        NOT YET BEEN ACKNOWLEDGED AS CANCELLED BY OctoPrint must still be
        "running". This prevents the cancelling printer from being rescheduled
        prematurely and from allows the job to be moved into the cancelled state
        by normal printer status inspection.

        """

        return self._conn.execute(
            """
            SELECT * FROM jobs
            WHERE started_at IS NOT NULL
                  AND printer_id IS NOT NULL
                  AND finished_at IS NULL
            """,
            [],
        ).fetchall()

    @requires_conn
    def list_cancelled_jobs(self):
        """Scheduler detail. List jobs which have been cancelled but are still 'live' (not finished)."""

        return self._conn.execute(
            """
            SELECT * FROM jobs
            WHERE finished_at IS NULL
                  AND cancelled_at IS NOT NULL
            """,
            [],
        ).fetchall()

    @fmap(one)
    @requires_conn
    def poll_job_queue(self):
        return self._conn.execute(
            """
            SELECT id
            FROM jobs
            WHERE started_at IS NULL
                  AND finished_at IS NULL
                  AND printer_id IS NULL
            ORDER BY priority DESC
            LIMIT 1
            """
        ).fetchone()

    @requires_conn
    def fetch_job(self, uid: int, jid: int) -> Optional[tuple]:
        return self._conn.execute(
            "SELECT * FROM jobs WHERE user_id = ? AND id = ?", [uid, jid]
        ).fetchone()

    @requires_conn
    def fetch_job_by_printer(self, pid: int) -> Optional[tuple]:
        """Find 'the' mapped incomplete job for a given printer."""

        return self._conn.execute(
            "SELECT * FROM jobs WHERE printer_id = ? AND finished_at IS NULL", [pid]
        ).fetchone()

    @requires_conn
    def assign_job(self, job_id: int, printer_id: int):
        return self._conn.execute(
            "UPDATE jobs SET printer_id = ?2 WHERE id = ?1", [job_id, printer_id]
        )

    @requires_conn
    def cancel_job(self, uid: int, job_id: int):
        return self._conn.execute(
            "UPDATE jobs SET cancelled_at = datetime('now') WHERE user_id = ?1 AND id = ?2", [uid, job_id]
        )

    @requires_conn
    def start_job(self, job_id: int):
        return self._conn.execute(
            "UPDATE jobs SET started_at = datetime('now') WHERE id = ?1", [job_id]
        )

    @requires_conn
    def finish_job(self, job_id: int, state: str, message: str = None):
        return self._conn.execute(
            "UPDATE jobs SET finished_at = datetime('now'), state = ?2, message = ?3 WHERE id = ?1",
            [job_id, state, message],
        )

    @requires_conn
    def delete_job(self, uid: int, jid: int):
        self._conn.execute("DELETE FROM jobs WHERE user_id = ? and id = ?", [uid, jid])
