from contextlib import closing

class ProgressStore:
    def _chunks(self, data, rows=10000):
        for i in range(0, len(data), rows):
            yield data[i:i+rows]

    def _execute(self, sql):
        with closing(self.connection.cursor()) as cursor:
            return [o for o in cursor.execute(sql)]

    def _commit(self):
        self.connection.commit()

    def __init__(self, db_location, bnumbers = None, table_name = 'progress'):
        self.db_location = db_location
        self.table_name = table_name

        self._bnumber_count = 0

        self.connection = sqlite3.connect(db_location)

        sql = (
            f"CREATE TABLE IF NOT EXISTS {self.table_name} "
            f"(bnumber text PRIMARY KEY, status text, notes text)"
        )

        self._execute(sql)
        self._commit()

        if (bnumbers):
            print(f"Initialising with reset!")
            self.reset(bnumbers)


    def count(self):
        sql = (
            f"SELECT COUNT(*) FROM {self.table_name}"
        )

        return self._execute(sql)[0][0]

    def reset(self, bnumbers):
        print(f"Resetting ProgressStore.")
        sql = f"DELETE FROM {self.table_name}"

        self._execute(sql)
        self._commit()

        self._bnumber_count = len(bnumbers)

        chunks = self._chunks(bnumbers)

        print(f"Loading {len(bnumbers)} bnumbers in chunks.")

        for chunk in chunks:
            self._execute('BEGIN TRANSACTION')

            for bnumber in chunk:
                sql = (
                    f"INSERT INTO {self.table_name} "
                    f"(bnumber) VALUES ('{bnumber}')"
                )
                self.cursor.execute(sql)

            self._execute('COMMIT')

        print("Done resetting.")

    def close(self):
        self.connection.close()

    def get_all(self):
        sql = (
            f"SELECT * FROM {self.table_name}"
        )

        with closing(self.connection.cursor()) as cursor:
            result = cursor.execute(sql)
            next_one = result.fetchone()

            while next_one is not None:
                yield next_one
                next_one = result.fetchone()

    def get_status(self, status):
        sql = (
            f"SELECT * FROM {self.table_name} "
            f"WHERE status = '{status}'"
        )

        return self._execute(sql)

    def get(self, bnumber):
        sql = (
            f"SELECT * FROM {self.table_name} "
            f"WHERE bnumber = '{bnumber}'"
        )

        return self._execute(sql)

    def update(self, bnumber, status, notes = None, commit = True):
        if(bnumber is None):
            raise Exception('Must include bnumber in status update!')

        if(status is None):
            raise Exception('Must include status in status update!')

        if(notes is None):
            notes = ""

        status = str(status).strip()

        sql = (
            f"UPDATE {self.table_name} SET "
            f"status = '{status}', notes = '{notes}' "
            f"WHERE bnumber = '{bnumber}'"
        )

        self._execute(sql)

        if commit:
            self._commit()

    def batch_update(self, status_updates):
        chunks = self._chunks(status_updates)

        print(f"Loading {len(status_updates)} status updates in chunks.")

        for chunk in chunks:
            self._execute('BEGIN TRANSACTION')

            for status_update in status_updates:
                self.update(
                    bnumber = status_update.get('bnumber'),
                    status = status_update.get('status'),
                    notes = status_update.get('notes'),
                    commit = False
                )

            self._execute('COMMIT')

        print(f"Done loading status updates.")