# Copyright (c) 2019 Mateusz Pawlik
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.

# ------------------------------------------------------------------------------
# DESCRIPTION
# ------------------------------------------------------------------------------
# This is an example implementation of the cookie jar exercise.
#
# Links to the documentation can be found at the end of this file and are
# referenced in the code with [X].
# ------------------------------------------------------------------------------

# Module for interaction with PostgreSQL [1-5].
import psycopg2

# Module for random number generation [6].
import random

# Module for concurrent execution of multiple processes [7].
from multiprocessing import Pool

# Module for processing command line parameters [8].
import argparse

# This module is used only to exit the program in case of a database error.
import sys

# Used to get a process id.
import os

# This method initialises the database, creates the cookie_jar table, and
# fills it in with cookies.
#
# num_cookies -- Number of cookies in the cookie jar.
def db_init(num_cookies):
    try:
        # Connect to database.
        db = psycopg2.connect(default_conn_string)
        # Open a cursor to perform database operations
        cur = db.cursor()

        # [NOTE] A command issued with cur.execute(...) starts a new transaction
        #        and lasts until db.commit() or db.rollback() are issued [3,5].

        # Reset the database by dropping table seats and recreating it.
        cur.execute("DROP TABLE IF EXISTS cookie_jar;")
        cur.execute("CREATE TABLE cookie_jar (\
                        cookie_id SERIAL,\
                        is_eaten integer DEFAULT 0\
                    );")
        # Make the changes to the database (DROP, CREATE) persistent.
        db.commit()

        # [NOTE] The message returned by the last command can be printed with
        #        'print cur.statusmessage'.

        # Put cookies into the jar.
        for _ in range(num_cookies):
            cur.execute("INSERT INTO cookie_jar DEFAULT VALUES")
        # Make the changes to the database (INSERT) persistent.
        db.commit()

        # Close the cursor.
        cur.close()
        # Close communication with the database (uncommitted changes will be
        # discarded as with issuing rollback).
        db.close()
    except Exception as e:
        print(type(e))
        print(e.args)
        sys.exit("Error occurred while initializing the database.")

# Eats a cookie.
# Returns the id of eaten cookie, or -1 if a cookie has not been eaten (for
# example, due to a database error).
#
# debug -- specifies if the error messages should be printed out
def monster_eat(default_conn_string, debug):
    # Get a process id to identify the monster.
    proc = os.getpid()
    # If anything goes wrong, this should be returned.
    failed_result = (proc, 'eating', -1)
    result = failed_result
    
    # Set cookie_is_overeaten.
    cookie_is_overeaten = False

    # Connect to database.
    try:
        db = psycopg2.connect(default_conn_string)
    except Exception as e:
        if debug:
            print(type(e))
            print(e.args)
        return failed_result

    # Logic of eating a cookie.
    try:
        
        # TODO: Implement eating a cookie here instead of the following 4 lines.
        cur = db.cursor()
        cur.execute("SELECT count(cookie_id) FROM cookie_jar;")
        cookies_in_jar = cur.fetchall()[0][0]
        result = (proc, 'eating', cookies_in_jar)
        
    except Exception as e:
        if debug:
            print(type(e))
            print(e.args)
        # Roll back the transaction on error, before the connection is closed.
        db.rollback()
        result = failed_result

    # Close communication with the database (uncommitted changes will be
    # discarded as with issuing a rollback).
    try:
        cur.close()
        db.close()
    except Exception as e:
        if debug:
            print(type(e))
            print(e.args)
        return failed_result

    return result

# Bakes a cookie.
# Returns the id of the baked cookie, or -1 if a cookie has not been baked (for
# example, due to a database error).
#
# debug -- specifies if the error messages should be printed out
def monster_bake(default_conn_string, debug):
    # Get a process id to identify the monster.
    proc = os.getpid()
    # If anything goes wrong, this should be returned.
    failed_result = (proc, 'baking', -1)
    result = failed_result

    # Connect to database.
    try:
        db = psycopg2.connect(default_conn_string)
    except Exception as e:
        if debug:
            print(type(e))
            print(e.args)
        return failed_result

    # Logic of baking a cookie.
    try:
        
        # TODO: Implement baking a cookie here instead of the following 4 lines.
        cur = db.cursor()
        cur.execute("SELECT count(cookie_id) FROM cookie_jar;")
        cookies_in_jar = cur.fetchall()[0][0]
        result = (proc, 'baking', cookies_in_jar)
        
    except Exception as e:
        if debug:
            print(type(e))
            print(e.args)
        # Roll back the transaction on error, before the connection is closed.
        db.rollback()
        result = failed_result

    # Close communication with the database (uncommitted changes will be
    # discarded as with issuing a rollback).
    try:
        cur.close()
        db.close()
    except Exception as e:
        if debug:
            print(type(e))
            print(e.args)
        return failed_result

    return result

# Monster, eat a cookie. If other monster steals your cookie and you find
# no more cookies in the jar, go to the kitchen and bake one.
# 
# For each cookie you bake, you get one more chance to eat one.
def monster_do(default_conn_string, debug):
    num_actions = 1
    results = []
    # Try to eat a cookie.
    r = monster_eat(default_conn_string, debug)
    eaten = r[2] > -1
    if eaten:
        results.append(r)
    else:
        # If you didn't eat a cookie, go and bake one.
        r = monster_bake(default_conn_string, debug)
        baked = r[2] > -1
        if baked:
            results.append(r)
            # Well done, you baked a cookie. You get one more try to eat a cookie.
            results.extend(monster_do(default_conn_string, debug))
    return results

# Collect some statistics about the cookie jar and print them to the command
# line in a JSON format.
def summarize_eating(args):
    result = ('{"m" : %d, "c" : %d, "e" : %d, ' % (args.num_monsters, args.init_cookies, args.num_cookies_to_eat))
    try:
        db = psycopg2.connect(args.conn_string)
        cur = db.cursor()
        cur.execute("SELECT count(cookie_id) FROM cookie_jar WHERE is_eaten=%s;", (0,))
        result += ('"cookies_in_jar" : %d,' % cur.fetchall()[0])
        cur.execute("SELECT count(cookie_id) FROM cookie_jar WHERE is_eaten>%s;", (0,))
        result += ('"cookies_eaten_from_jar" : %d,' % cur.fetchall()[0])
        cur.execute("SELECT count(cookie_id) FROM cookie_jar WHERE is_eaten>%s;", (1,))
        result += ('"cookies_eaten_more_than_once" : %d,' % cur.fetchall()[0])
        cur.execute("SELECT sum(is_eaten) FROM cookie_jar;")
        result += ('"total_cookies_eaten" : %d,' % cur.fetchall()[0])
        cur.execute("SELECT count(cookie_id) FROM cookie_jar;")
        result += ('"cookies_baked" : %d' % (int(cur.fetchall()[0][0]) - init_cookies))
        db.commit()
        cur.close()
        db.close()
    except Exception as e:
        if debug:
            print(type(e))
            print(e.args)
        sys.exit("Error occurred while collecting statistics.")
    result += "}"
    return result

# Verifies the number of monsters.
#
# x -- Number of monsters to verify.
def number_monsters_type(x):
    x = int(x)
    if x > 20:
        raise argparse.ArgumentTypeError("Maximum number of monsters is 20.")
    return x

# ------------------------------------------------------------------------------
# EXECUTE SIMULATOR
# ------------------------------------------------------------------------------
if __name__ == '__main__':
    # Processing command line parameters
    parser = argparse.ArgumentParser(
        description='Advanced Databases - Concurrent Cookie Jar. \
                     Description can be found in the source code.')
    parser.add_argument(
        '--initial-cookies', '-c',
        type=int,
        dest='init_cookies',
        help='Number of cookies (default: 10).',
        default=10)
    parser.add_argument(
        '--cookies-to-eat', '-e',
        type=int,
        dest='num_cookies_to_eat',
        help='Number of cookies to eat (default: 20).',
        default=20)
    parser.add_argument(
        '--monsters', '-m',
        type=number_monsters_type,
        dest='num_monsters',
        help='Number of monsters (default: 2).',
        default=2)
    parser.add_argument(
        '--connection-string',
        type=str,
        dest='conn_string',
        action='store',
        help="Database connection string (default: empty string)",
        default="")
    parser.add_argument(
        '-d', '--debug',
        dest='debug',
        action='store_true',
        help="Debug mode - prints database errors if any.")
    args = parser.parse_args()
    # Set the DB connection string.
    default_conn_string = args.conn_string
    # Set the number of cookies.
    init_cookies = args.init_cookies
    num_cookies_to_eat = args.num_cookies_to_eat
    # Set the number of monsters.
    num_monsters = args.num_monsters
    # Set the debug flag.
    debug=args.debug

    # Initialize the database.
    db_init(init_cookies)
    
    # Collect arguments for all `monster_do` executions.
    do_args = [(default_conn_string, debug) for _ in range(num_cookies_to_eat)]
    # Create a pool of monsters.
    with Pool(num_monsters) as pool:
        # Execute transactions concurrently.
        # Collect the results of concurrent execution of `monster_do`
        # in `multiple_results` list.
        multiple_results = pool.starmap(monster_do, do_args)

    # If in debug mode, print the result of each execution of `monster_do`.
    if debug:
        print("Order of monsters doing things:")
        for r in multiple_results:
            t = ''
            for r_elem in r:
                if r_elem[2] == -1:
                    print(t, "Monster %d didn't %s cookie." % (r_elem[0], ('eat' if r_elem[1] == 'eating' else 'bake')))
                else:
                    print(t, "Monster %d %s cookie %d." % (r_elem[0], ('ate' if r_elem[1] == 'eating' else 'baked'), r_elem[2]))
                t += '\t'

    # Report on cookies eaten.
    print(summarize_eating(args))
# ------------------------------------------------------------------------------

# ------------------------------------------------------------------------------
# REFERENCES
# ------------------------------------------------------------------------------
# psycopg2:
# [1] basic usage
#     http://initd.org/psycopg/docs/usage.html
# [2] transaction control
#     http://initd.org/psycopg/docs/usage.html#transactions-control
# [3] connection - transaction handling with commit() and rollback()
#     http://initd.org/psycopg/docs/connection.html
# [4] cursor - for query execution and fetching the results
#     http://initd.org/psycopg/docs/cursor.html
# [5] examples of transaction handling with comparisons to PostgreSQL
#     https://wiki.python.org/moin/UsingDbApiWithPostgres#transactions
# Python:
# [6] random - random number generation
#     https://docs.python.org/3/library/random.html
# [7] multiprocessing - executing methods as separate processes with examples
#     https://docs.python.org/3/library/multiprocessing.html
#     https://docs.python.org/3/library/multiprocessing.html#examples
# [8] argparse - parsing command line parameters with tutorial
#     https://docs.python.org/3/library/argparse.html
#     https://docs.python.org/3/howto/argparse.html
# ------------------------------------------------------------------------------