################################
#### Name: Yuanfeng Wu
#### Uniqname: yuanfenw
################################

import sqlite3
import plotly.graph_objects as go 

# proj3_choc.py
# You can change anything in this file you want as long as you pass the tests
# and meet the project requirements! You will need to implement several new
# functions.

# Part 1: Read data from a database called choc.db
DBNAME = 'choc.sqlite'

# Part 1: Implement logic to process user commands
def update_database(DBNAME):
    '''update the database to add CompanyLocation and BroadBeanOrigin
        
    create a new table bars_countries
    copy values of table bars to bars_countries
    add CompanyLocation and BroadBeanOrigin to bars_countries

    Parameters  
    ----------
    DBNAME: str
        the path to database file

    Returns
    -------
    None
    '''
    conn = sqlite3.connect(DBNAME)
    cursor = conn.cursor()
    query = '''
        SELECT EnglishName
        FROM Countries
        JOIN Bars
        WHERE Bars.CompanyLocationId = Countries.Id
    '''
    res = cursor.execute(query)
    company_names = [r[0] for r in res]

    query = '''
        SELECT EnglishName
        FROM Countries
        JOIN Bars
        WHERE Bars.BroadBeanOriginId = Countries.Id
    '''
    res = cursor.execute(query)
    beanorigin_names = [r[0] for r in res]

    drop = '''
        DROP TABLE IF EXISTS bars_countries
    '''
    create = '''
        CREATE TABLE 'bars_countries' (
            'Id' INTEGER PRIMARY KEY AUTOINCREMENT,
            'Company' TEXT NOT NULL,
            'SpecificBeanBarName' TEXT NOT NULL,
            'REF' TEXT NOT NULL,
            'ReviewDate' TEXT NOT NULL,
            'CocoaPercent' REAL,
            'CompanyLocationId' INTEGER,
            'Rating' REAL,
            'BeanType' TEXT NOT NULL,
            'BroadBeanOriginId' INTEGER,
            'CompanyLocation' TEXT,
            'BroadBeanOrigin' TEXT
        );
    '''
    cursor.execute(drop)
    cursor.execute(create)
    conn.commit()

    query = '''
        SELECT * FROM bars
    '''
    res = cursor.execute(query).fetchall()
    for i in range(len(company_names)):
        data = list(res[i]) + [company_names[i]] + [beanorigin_names[i]]
        query = '''
            INSERT INTO bars_countries 
            VALUES(?,?,?,?,?,?,?,?,?,?,?,?)
        '''
        cursor.execute(query, data)
        conn.commit()
       
def comm_bars(command):
    '''Implement bars query
        
    Lists chocolate bars, according to the specified parameters.

    Parameters  
    ----------
    command: str
        the bar query

    Returns
    -------
    result: 
        list: a list of 6-tuple records when work properly
            (Specific Bean Bar Name, Company Name, Company Location, 
            Rating, CocoaPercent, Broad Bean Origin)
        str: error notification
    '''
    conn = sqlite3.connect(DBNAME)
    cur = conn.cursor()
    command_lst = command.strip().split()
    params = command_lst[1:]
    # default setting
    region = None
    region_name = None
    order = 'ratings'
    limit = 'top'
    num = 10
    sell_source = None
    plot = 0
    for p in params:
        p_list = p.split('=')
        p_key = p_list[0]
        if p_key in ['country', 'region']:
            region = p_key
            try:
                region_name = p_list[1]
            except:
                return 'Missing value for parameters. Please try again.'
        elif p_key in ['sell', 'source']:
            sell_source = p_key
        elif p_key in ['ratings', 'cocoa']:
            order = p_key
        elif p_key in ['top', 'bottom']:
            limit = p_key
        elif p_key.isnumeric():
            num = int(p_key)
            if num <= 0:
                return 'Invalid limit number "{}". Please try again.'.format(num)
        elif p_key == 'barplot':
            plot = 1
        else:
            return 'Invalid parameter "{}". Please try again.'.format(p_key)

    if order == 'ratings':
        order_column = 'Rating'
    else:
        order_column = 'CocoaPercent'

    if limit == 'top':
        order_direction = 'DESC'
    else:
        order_direction = 'ASC'

    if sell_source == 'sell':
        if region == 'country':
            condition = '''
                JOIN Countries 
                    ON bars_countries.CompanyLocationId = Countries.Id
                WHERE Alpha2 = '{}'
            '''.format(region_name)
        elif region == 'region':
            condition = '''
                JOIN Countries 
                    ON bars_countries.CompanyLocationId = Countries.Id
                WHERE Region = '{}'
            '''.format(region_name)
    elif sell_source == 'source':
        if region == 'country':
            condition = '''
                JOIN Countries 
                    ON bars_countries.BroadBeanOriginId = Countries.Id
                WHERE Alpha2 = '{}'
            '''.format(region_name)
        elif region == 'region':
            condition = '''
                JOIN Countries 
                    ON bars_countries.BroadBeanOriginId = Countries.Id
                WHERE Region = '{}'
            '''.format(region_name)
    else:
        condition = ''

    statement = '''
        SELECT SpecificBeanBarName, Company, CompanyLocation, Rating, CocoaPercent, BroadBeanOrigin
        FROM bars_countries
        {}
        ORDER BY {} {}
        LIMIT {}
    '''.format(condition, order_column, order_direction, num)
    cur.execute(statement)
    return_lst = cur.fetchall()
    result = []
    for r in return_lst:
        result.append(r)
    if plot == 1:
        xvals = [r[0] for r in result]
        if order == 'ratings':
            yvals = [r[3] for r in result]
        else:
            yvals = [r[4] for r in result]
        bar_data = go.Bar(x=xvals, y=yvals)
        basic_layout = go.Layout(title="A Bar Graph")
        fig = go.Figure(data=bar_data, layout=basic_layout)
        fig.show()

    return result

def comm_companies(command):
    '''Implement companies query
        
    Lists chocolate bars sellers according to the specified parameters.
    Only companies that sell more than 4 bars are listed in results.

    Parameters  
    ----------
    command: str
        the companies query

    Returns
    -------
    result: 
        list: a list of 3-tuple records when work properly
            (Company Name, Company Location, <agg: avg rating, avg cocoa, or number_of_bars>)
        str: error notification
    '''
    conn = sqlite3.connect(DBNAME)
    cur = conn.cursor()
    command_lst = command.strip().split()
    params = command_lst[1:]
    # default setting
    region = None
    region_name = None
    order = 'ratings'
    limit = 'top'
    num = 10
    plot = 0
    for p in params:
        p_list = p.split('=')
        p_key = p_list[0]
        if p_key in ['country', 'region']:
            region = p_key
            try:
                region_name = p_list[1]
            except:
                return 'Missing value for parameters. Please try again.'
        elif p_key in ['ratings', 'cocoa', 'number_of_bars']:
            order = p_key
        elif p_key in ['top', 'bottom']:
            limit = p_key
        elif p_key.isnumeric():
            num = int(p_key)
            if num <= 0:
                return 'Invalid limit number "{}". Please try again.'.format(num)
        elif p_key == 'barplot':
            plot = 1
        else:
            return 'Invalid parameter "{}". Please try again.'.format(p_key)
    # order statement
    if order == 'ratings':
        order_criteria = 'AVG(Rating)'
    elif order == 'cocoa':
        order_criteria = 'AVG(CocoaPercent)'
    else:
        order_criteria = 'COUNT(*)'

    if limit == 'top':
        order_direction = 'DESC'
    else:
        order_direction = 'ASC'

    # having_and statement
    if region == 'country':
        having = 'AND Alpha2 = \'{}\''.format(region_name)
    elif region == 'region':
        having = 'AND Region = \'{}\''.format(region_name)
    else:
        having = ''

    statement = '''
        SELECT Company, CompanyLocation, {}
        FROM bars_countries
            JOIN Countries
                ON bars_countries.CompanyLocationId = Countries.Id
        GROUP BY Company
        HAVING COUNT(*) > 4 {}
        ORDER BY {} {}
        LIMIT {}
    '''.format(order_criteria, having, order_criteria, order_direction, num)
    cur.execute(statement)
    return_lst = cur.fetchall()
    result = []
    for r in return_lst:
        result.append(r)
    if plot == 1:
        xvals = [r[0] for r in result]
        yvals = [r[2] for r in result]
        bar_data = go.Bar(x=xvals, y=yvals)
        basic_layout = go.Layout(title="A Bar Graph")
        fig = go.Figure(data=bar_data, layout=basic_layout)
        fig.show()
    return result

def comm_countries(command):
    '''Implement countries query
        
    Lists countries according to specified parameters.
    Only countries that sell/source more than 4 bars are listed in results.


    Parameters  
    ----------
    command: str
        the countries query

    Returns
    -------
    result: 
        list: a list of 3-tuple records when work properly
            (Country, Region, <agg: avg rating, avg cocoa, or number_of_bars>)
        str: error notification
    '''
    conn = sqlite3.connect(DBNAME)
    cur = conn.cursor()
    command_lst = command.strip().split()
    params = command_lst[1:]
    # default setting
    region = None
    region_name = None
    sell_source = 'sell'
    order = 'ratings'
    limit = 'top'
    num = 10
    plot = 0
    for p in params:
        p_list = p.split('=')
        p_key = p_list[0]
        if p_key in ['region']:
            region = p_key
            try:
                region_name = p_list[1]
            except:
                return 'Missing value for parameters. Please try again.'
        elif p_key in ['sell', 'source']:
            sell_source = p_key
        elif p_key in ['ratings', 'cocoa', 'number_of_bars']:
            order = p_key
        elif p_key in ['top', 'bottom']:
            limit = p_key
        elif p_key.isnumeric():
            num = int(p_key)
            if num <= 0:
                return 'Invalid limit number "{}". Please try again.'.format(num)
        elif p_key == 'barplot':
            plot = 1
        else:
            return 'Invalid parameter "{}". Please try again.'.format(p_key)

    # order statement
    if order == 'ratings':
        order_criteria = 'AVG(Rating)'
    elif order == 'cocoa':
        order_criteria = 'AVG(CocoaPercent)'
    else:
        order_criteria = 'COUNT(*)'

    if limit == 'top':
        order_direction = 'DESC'
    else:
        order_direction = 'ASC'

    # having_and statement
    if region == 'region':
        having = 'AND Region = \'{}\''.format(region_name)
    else:
        having = ''

    # join statement
    if sell_source == 'sell':
        join = 'CompanyLocationId'
    else:
        join = 'BroadBeanOriginId'

    statement = '''
        SELECT Countries.EnglishName, Countries.Region, {}
        FROM bars_countries
            JOIN Countries
                ON bars_countries.{} = Countries.Id
        GROUP BY Countries.Alpha2
        HAVING COUNT(*) > 4 {}
        ORDER BY {} {}
        LIMIT {}
    '''.format(order_criteria, join, having, order_criteria, order_direction, num)

    cur.execute(statement)
    return_lst = cur.fetchall()
    result = []
    for r in return_lst:
        result.append(r)
    if plot == 1:
        xvals = [r[0] for r in result]
        yvals = [r[2] for r in result]
        bar_data = go.Bar(x=xvals, y=yvals)
        basic_layout = go.Layout(title="A Bar Graph")
        fig = go.Figure(data=bar_data, layout=basic_layout)
        fig.show()
    return result

def comm_regions(command):
    '''Implement regions query
        
    Lists regions according to specified parameters.
    Only regions that sell/source more than 4 bars are listed in results.


    Parameters  
    ----------
    command: str
        the regions query

    Returns
    -------
    result: 
        list: a list of 2-tuple records when work properly
            (Region, <agg: avg rating, avg cocoa, or number_of_bars>)
        str: error notification
    '''
    conn = sqlite3.connect(DBNAME)
    cur = conn.cursor()
    command_lst = command.strip().split()
    params = command_lst[1:]
    # default setting
    sell_source = 'sell'
    order = 'ratings'
    limit = 'top'
    num = 10
    plot = 0
    for p in params:
        p_list = p.split('=')
        p_key = p_list[0]
        if p_key in ['sell', 'source']:
            sell_source = p_key
        elif p_key in ['ratings', 'cocoa', 'number_of_bars']:
            order = p_key
        elif p_key in ['top', 'bottom']:
            limit = p_key
        elif p_key.isnumeric():
            num = int(p_key)
            if num <= 0:
                return 'Invalid limit number "{}". Please try again.'.format(num)
        elif p_key == 'barplot':
            plot = 1
        else:
            return 'Invalid parameter "{}". Please try again.'.format(p_key)

    # order statement
    if order == 'ratings':
        order_criteria = 'AVG(Rating)'
    elif order == 'cocoa':
        order_criteria = 'AVG(CocoaPercent)'
    else:
        order_criteria = 'COUNT(*)'

    if limit == 'top':
        order_direction = 'DESC'
    else:
        order_direction = 'ASC'

    # join statement
    if sell_source == 'sell':
        join = 'CompanyLocationId'
    else:
        join = 'BroadBeanOriginId'

    statement = '''
        SELECT Region, {}
        FROM Bars
            JOIN Countries
                ON Bars.{} = Countries.Id
        GROUP BY Region
        HAVING COUNT(*) > 4
        ORDER BY {} {}
        LIMIT {}
    '''.format(order_criteria, join, order_criteria, order_direction, num)

    cur.execute(statement)
    return_lst = cur.fetchall()
    result = []
    for r in return_lst:
        result.append(r)
    if plot == 1:
        xvals = [r[0] for r in result]
        yvals = [r[1] for r in result]
        bar_data = go.Bar(x=xvals, y=yvals)
        basic_layout = go.Layout(title="A Bar Graph")
        fig = go.Figure(data=bar_data, layout=basic_layout)
        fig.show()
    return result

def process_command(command):
    '''Implement the query
    
    Takes a command and returns records that match the query

    Parameters  
    ----------
    command: str
        the query

    Returns
    -------
    result: 
        list: a list of <integer>-tuple records when work properly
        str: error notification
    '''
    command_list = command.strip().split(' ')
    if len(command_list)==0:
        res = 'Invalid command. Please try again.'
    else:
        head = command_list[0]
        if head == 'bars':
            res = comm_bars(command)
        elif head == 'companies':
            res = comm_companies(command)
        elif head == 'countries':
            res = comm_countries(command)
        elif head == 'regions':
            res = comm_regions(command)
        else:
            res = 'Invalid command "{}". Please try again.'.format(head)
    return res


def load_help_text():
    '''load help text

    Parameters  
    ----------
    None

    Returns
    -------
    res: str
        the help documentation
    '''
    with open('Proj3Help.txt') as f:
        return f.read()

# Part 2 & 3: Implement interactive prompt and plotting. We've started for you!
def interactive_prompt():
    '''allow user to input commands and receive results

    allow a user to interactively input commands 
    and to nicely format the results for presentation.

    Parameters  
    ----------
    None

    Returns
    -------
    None
    '''
    update_database(DBNAME)
    help_text = load_help_text()
    response = ''
    while True:
        response = input('Enter a command: ')
        if response == 'exit':
            break
        if response == 'help':
            print(help_text)
            continue
        else:
            return_lst = process_command(response)
            if isinstance(return_lst, str):
                print(return_lst)
            elif return_lst == []:
                print("No records found.")
            else:
                main = response.split()[0]
                if main in ['companies', 'countries', 'regions']:
                    agg = 'ratings'
                    try:
                        params = response.split()[1:]
                        for p in params:
                            p_list = p.split('=')
                            p_key = p_list[0]
                            if p_key in ['ratings', 'cocoa', 'number_of_bars']:
                                agg = p_key
                    except:
                        pass
                    # <agg> part in output string
                    if agg == 'ratings':
                        agg_s = '1.1f'
                    elif agg == 'cocoa':
                        agg_s = '4.0%'
                    else:
                        agg_s = '4d'

                for row in return_lst:
                    row_lst = list(row)

                    # limit string length
                    for i, e in enumerate(row_lst):
                        if isinstance(e, str) and len(e) > 15:
                            row_lst[i] = e[:12] + '...'

                    if main == 'bars':
                        print('{0:15s} {1:15s} {2:15s} {3:1.1f} {4:4.0%} {5:15s}'.format(*row_lst))
                    elif main in ['companies', 'countries']:
                        output = '{0:15s} {1:15s} {2:' + agg_s + '}'
                        print(output.format(*row_lst))
                    else:
                        output = '{0:15s} {1:' + agg_s + '}'
                        print(output.format(*row_lst))
    print('Bye')

# Make sure nothing runs or prints out when this file is run as a module/library
if __name__=="__main__":
    interactive_prompt()