#!/usr/bin/env python3

import argparse
import plotly.express as px
import pandas as pd
#import datetime
import sqlite3
import time
import re
import sys
from operator import itemgetter
import shutil

from os import listdir
from dash import Dash, dcc, html, dash_table, Input, Output, State
from dash.exceptions import PreventUpdate
import dash_bootstrap_components as dbc

AS = {}
AS['directory'] = None
AS['database'] = None
AS['argv'] = None
AS['start_time'] = None

def update_legend_labels(fig,labels):
    for entry in fig.data:
        if entry['name'] in labels:
            entry['name'] = labels[entry['name']]

def get_number_of_experiments(directory, database):
    conn = sqlite3.connect(f"file:{directory}/{database}?mode=ro", uri=True)
    cursor = conn.cursor()
    query = "SELECT COUNT(*) FROM experiments"
    cursor.execute(query)
    result = cursor.fetchone()
    row_count = result[0]
    conn.close()
    return row_count

def get_start_time(directory, database):
    conn = sqlite3.connect(f"file:{directory}/{database}?mode=ro", uri=True)
    cursor = conn.cursor()
    query = "SELECT stime_seconds FROM metadata"
    cursor.execute(query)
    result = cursor.fetchone()
    seconds = result[0]
    conn.close()
    return seconds

def update_metadata(directory, database):
    AS['directory'] = directory
    AS['database'] = database
    conn = sqlite3.connect(f"file:{directory}/{database}?mode=ro", uri=True)
    cursor = conn.cursor()
    query = "SELECT * FROM metadata"
    cursor.execute(query)
    result = cursor.fetchone()
    AS['start_time'] = result[0]
    conn.close()
    try:
        conn = sqlite3.connect(f"file:{directory}/{database}?mode=ro", uri=True)
        cursor = conn.cursor()    
        query = "SELECT argv FROM metadata"
        cursor.execute(query)
        result = cursor.fetchone()
        AS['argv'] = result[0]
        conn.close()
    except Exception as _:
        AS['argv'] = 'Missing from database'

def get_databases(directory):
    # get databases
    databases = []
    for file in listdir(directory):
        if re.search('^.*\\.sqlite$',file):
            databases.append(file)
    databases.sort(reverse=True)
    # add number of experiments
    databases_new = []
    for index in range(len(databases)):
        databases_new.append('%s (%d)' %(databases[index], get_number_of_experiments(directory, databases[index])))
    return databases_new

def run(directory, ip="127.0.0.1", port=8080, x_axis="delay", y_axis="length", aspect_ratio=0, auto_update_interval=0, debug=False):
    DATABASE_DIRECTORY = directory
    app = Dash(__name__, external_stylesheets=[dbc.themes.LUX])
    app.css.config.serve_locally = True
    app.scripts.config.serve_locally = True
    app.layout = html.Div([
            html.Div([
                html.H4('Fault Injection Analysis'),
            ],style={'width':'80%','border-style':'none','margin':'0 auto'}),               
            html.Div([
                html.Button("Update", id='update-button', n_clicks=0, style={'width':'20%'}),
                html.Datalist(id="examples", children=[
                    html.Option(value="match_string(response, 'ets')"),
                    html.Option(value="match_hex(response, '661b')"),
                    html.Option(value="color = 'G'"),
                    html.Option(value=f"length > 100"),
                    html.Option(value=f"delay > 100"),
                    html.Option(value=f"voltage = 400"),
                ]),
                dcc.Input(id='query-input', type="text", list='examples', style={'width':'80%','display': 'inline-block'}, placeholder="SELECT * FROM experiments WHERE"),
                dcc.Dropdown(id='graph-dropdown', style={'width':'100%'}),
                html.Center([
                    dcc.Graph(id='graph', style={'width':'80%'}),
                ]),
                html.P('re.search(*, response)'),
                dcc.Input(id='recolor-green', type="text", placeholder="green", style={'width':'15%'}),
                dcc.Input(id='recolor-green-label', type="text", placeholder="green-label", style={'width':'10%'}),
                dcc.Input(id='recolor-yellow', type="text", placeholder="yellow", style={'width':'15%'}),
                dcc.Input(id='recolor-yellow-label', type="text", placeholder="yellow-label", style={'width':'10%'}),
                dcc.Input(id='recolor-magenta', type="text", placeholder="magenta", style={'width':'15%'}),
                dcc.Input(id='recolor-magenta-label', type="text", placeholder="magenta-label", style={'width':'10%'}),
                dcc.Input(id='recolor-orange', type="text", placeholder="orange", style={'width':'15%'}),
                dcc.Input(id='recolor-orange-label', type="text", placeholder="orange-label", style={'width':'10%'}),
                html.Br(),
                dcc.Input(id='recolor-cyan', type="text", placeholder="cyan", style={'width':'15%'}),
                dcc.Input(id='recolor-cyan-label', type="text", placeholder="cyan-label", style={'width':'10%'}),
                dcc.Input(id='recolor-blue', type="text", placeholder="blue", style={'width':'15%'}),
                dcc.Input(id='recolor-blue-label', type="text", placeholder="blue-label", style={'width':'10%'}),
                dcc.Input(id='recolor-black', type="text", placeholder="black", style={'width':'15%'}),
                dcc.Input(id='recolor-black-label', type="text", placeholder="black-label", style={'width':'10%'}),
                dcc.Input(id='recolor-red', type="text", placeholder="red", style={'width':'15%'}),
                dcc.Input(id='recolor-red-label', type="text", placeholder="red-label", style={'width':'10%'}),

                html.Br(),
                html.Label('Combine data?'),
                dcc.RadioItems(id='combine-data', options=['Yes', 'No'], value='Yes', inline=True),
                html.Div(id='data',style={'width':'100%', 'height':'80%', 'border-style':'none'}),
                html.Label('Arguments:'),
                html.Label('lalala', id='argv'),

                # timer for update_graph callback
                dcc.Interval(
                    id="interval-component",
                    interval=auto_update_interval * 1000,
                    n_intervals=0,
                    disabled=(auto_update_interval == 0)
                )
            ],style={'width':'80%','border-style':'none','margin':'0 auto'}),
        ],style={'width':'100%', 'border-style':'none', 'margin-top':'100px','margin-bottom':'100px'})

    # callback for database list
    @app.callback(
        Output("argv", "children"),
        Input('update-button', 'n_clicks')
    )
    def update_argv(n_clicks):
        return AS['argv']

    # callback for database list
    @app.callback(
        Output("graph-dropdown", "options"),
        Input('update-button', 'n_clicks')
    )
    def update_dropdown(n_clicks):
        return get_databases(DATABASE_DIRECTORY)

    def percentage(val=None,total=None):
        return "{:.1%}".format(0.1234)

    # new function for sqlite3 query
    def match_string(response, token):
        if token.encode(errors='strict') in response:
            return True
        else:
            return False

    # new function for sqlite3 query
    def match_hex(response, token):
        if bytes.fromhex(token) in response:
            return True
        else:
            return False

    def recolor(regex, response):
        if regex is None or len(regex) == 0:
            return False
        elif re.search(regex.encode(), response):
            return True
        else:
            return False

    # callback for database selection  
    @app.callback(
        Output('graph','figure'),
        Output('data','children'),
        Input('update-button', 'n_clicks'),
        Input('graph-dropdown', 'value'),
        Input('interval-component', 'n_intervals'),
        State('query-input', 'value'),
        State('recolor-green', 'value'),
        State('recolor-green-label', 'value'),
        State('recolor-yellow', 'value'),
        State('recolor-yellow-label', 'value'),
        State('recolor-magenta', 'value'),
        State('recolor-magenta-label', 'value'),
        State('recolor-orange', 'value'),
        State('recolor-orange-label', 'value'),
        State('recolor-cyan', 'value'),
        State('recolor-cyan-label', 'value'),
        State('recolor-blue', 'value'),
        State('recolor-blue-label', 'value'),
        State('recolor-black', 'value'),
        State('recolor-black-label', 'value'),
        State('recolor-red', 'value'),
        State('recolor-red-label', 'value'),
        State('combine-data', 'value'),
    )
    def update_graph(nr_of_clicks, database, interval, query, green, greenlabel, yellow, yellowlabel, magenta, magentalabel, orange, orangelabel, cyan, cyanlabel, blue, bluelabel, black, blacklabel, red, redlabel, combine):
        if debug:
            now = round(time.time() * 1000)
        
        # update databse 
        if not database:
            raise PreventUpdate
        
        database = database.split(' ')[0]

        # copy database to /tmp and opening it
        print(f"Copying {database} to /tmp and opening from there")
        shutil.copyfile(f"{DATABASE_DIRECTORY}/{database}", f"/tmp/{database}")
        con = sqlite3.connect(f"file:/tmp/{database}?mode=ro", uri=True)

        con.create_function('match_string', 2, match_string)
        con.create_function('match_hex', 2, match_hex)

        # updating metadata from database
        update_metadata(DATABASE_DIRECTORY, database)

        # perform the query based on the query extension
        if query is not None and query != '':
            query = f'SELECT * FROM experiments WHERE %s;' %(query)
        else:
            query = 'SELECT * FROM experiments;'

        # read stuff from database
        df = pd.read_sql(query, con)
        con.close()

        # recolor if needed
        records = df.to_dict('records')
        for record in records:
            if recolor(green, record['response']):
                record['color'] = 'G'
            elif recolor(yellow, record['response']):
                record['color'] = 'Y'
            elif recolor(magenta, record['response']):
                record['color'] = 'M'
            elif recolor(orange, record['response']):
                record['color'] = 'O'
            elif recolor(cyan, record['response']):
                record['color'] = 'C'
            elif recolor(blue, record['response']):
                record['color'] = 'B'
            elif recolor(black, record['response']):
                record['color'] = 'Z'
            elif recolor(red, record['response']):
                record['color'] = 'R'
        
        # create new DataFrame of recolored data
        df = pd.DataFrame.from_dict(records)

        # get amount of experiments
        nr_of_current_experiments = len(df) 

        # output plot
        fig = px.scatter(
            df,
            x = x_axis,
            y = y_axis,
            render_mode = "webgl",
            color = "color", 
            labels = {"color":f"Classification ({nr_of_current_experiments:,})",y_axis:y_axis, x_axis:x_axis},
            color_discrete_map = { 
                "G": "green",
                "Y": "yellow", 
                "M" : "magenta", 
                "O": "orange",
                "C": "cyan",
                "B": "blue",
                "Z": "black",
                "R": "red", 
            },
            category_orders = {
                "color" : ["G","Y","M","O","C","B","R"]
            },
        )

        # compute elapsed time
        # elapsed_time_in_seconds = int(time.time()-get_start_time(DATABASE_DIRECTORY, database))
        # elapsed_time = datetime.timedelta(seconds=elapsed_time_in_seconds)
        # nr_of_experiments_per_second = nr_of_current_experiments // elapsed_time_in_seconds
        
        # update title of graph
        # fig.update_layout(title_text=f"{database[:-7]} (running for {elapsed_time} @ {nr_of_experiments_per_second} per second)\n", title_x=0.5)
        if aspect_ratio != 0:
            fig.update_layout(title_text=f"{database[:-7]}", title_x=0.5, yaxis=dict(scaleanchor="x", scaleratio=aspect_ratio))
        else:
            fig.update_layout(title_text=f"{database[:-7]}", title_x=0.5)

        # update labels in legenda
        def make_label(color, label, df):
            count = len(df.query(f"color == '{color}'"))
            if count > 0:
                percentage = "{:.1%}".format(count/len(df))
            else:
                percentage = "{:.1%}".format(0)
            return { color: f'{label} ( {count} / {percentage} )'}

        labels = {}
        labels.update(make_label('G', greenlabel, df))
        labels.update(make_label('Y', yellowlabel, df))
        labels.update(make_label('M', magentalabel, df))
        labels.update(make_label('O', orangelabel, df))
        labels.update(make_label('C', cyanlabel, df))
        labels.update(make_label('B', bluelabel, df))
        labels.update(make_label('Z', blacklabel, df))
        labels.update(make_label('R', redlabel, df))
        update_legend_labels(fig, labels)

        # output data
        records = df.to_dict('records')

        y_min = y_max = x_min = x_max = 0

        # TODO: rewrite this code
        if combine == 'Yes':

            combined_records = []
            for record in records:

                # decode response to make sure it's compatible with json
                response_hex = record['response'].hex(' ')
                response = record['response'].decode('utf-8', errors='replace')

                # search for response that's already there
                for index in range(len(combined_records)):
                    if combined_records[index]['response'] == response:
                        combined_records[index]['amount'] += 1
                        
                        if record[y_axis] < combined_records[index]['yMin']:
                            combined_records[index]['yMin'] = record[y_axis]

                        if record[y_axis] > combined_records[index]['yMax']:
                            combined_records[index]['yMax'] = record[y_axis]

                        if record[x_axis] < combined_records[index]['xMin']:
                            combined_records[index]['xMin'] = record[x_axis]

                        if record[x_axis] > combined_records[index]['xMax']:
                            combined_records[index]['xMax'] = record[x_axis]
                        break
                else:
      

                    y_min = y_max = record[y_axis]
                    x_min = x_max = record[x_axis]

                    combined_records.append({'amount': 1, 'color': record['color'], 'yMin': y_min, 'yMax': y_max, 'xMin': x_min, 'xMax': x_max, 'response': response, 'response_hex': response_hex})

            # sort new list based on occurrences 
            combined_records = sorted(combined_records, key=itemgetter('amount'), reverse=True)

            columns = ['amount', 'color', 'yMin', 'yMax', 'xMin', 'xMax', 'response', 'response_hex' ]

            cell_style = [
                {'if': {'column_id': 'amount'},     'textAlign': 'center','width':'100px'},
                {'if': {'column_id': 'color'},      'textAlign': 'center','width':'100px'},
                {'if': {'column_id': 'yMin'},   'textAlign': 'center','width':'100px'},
                {'if': {'column_id': 'yMax'},   'textAlign': 'center','width':'100px'},
                {'if': {'column_id': 'xMin'},  'textAlign': 'center','width':'100px'},
                {'if': {'column_id': 'xMax'},  'textAlign': 'center','width':'100px'},
                {'if': {'column_id': 'response'},   'textAlign': 'left'},
                {'if': {'column_id': 'response_hex'},     'textAlign': 'left'}
            ]

            nr_of_fixed_columns = 6

            records = combined_records
 
        else:
            columns = ['id', 'color', y_axis, x_axis, 'rlen', 'response','response_hex']

            all_records = []

            for record in records:
                new_record = {}
                new_record['id'] = record['id']
                new_record[y_axis] = record[y_axis]
                new_record[x_axis] = record[x_axis]
                new_record['color'] = record['color']
                new_record['rlen'] = len(record['response'])
                new_record['response'] = record['response'].decode('utf-8', errors='replace')
                new_record['response_hex'] = record['response'].hex(' ')
                all_records.append(new_record)

            cell_style = [
                {'if': {'column_id': 'id'},         'textAlign': 'center','width':'100px', 'minWidth':'100px'},
                {'if': {'column_id': y_axis},      'textAlign': 'center','width':'100px', 'minWidth':'100px'},
                {'if': {'column_id': x_axis},     'textAlign': 'center','width':'100px', 'minWidth':'100px'},
                {'if': {'column_id': 'color'},      'textAlign': 'center','width':'100px', 'minWidth':'100px'},
                {'if': {'column_id': 'rlen'},     'textAlign': 'center','width':'100px', 'minWidth':'100px'},
                {'if': {'column_id': 'response'},   'textAlign': 'left'},
                {'if': {'column_id': 'response_hex'},     'textAlign': 'left'},
            ]

            nr_of_fixed_columns = 5
            records = all_records

        data_style = [
            {'if': {'filter_query': '{color} = G'},'backgroundColor': 'green','color': 'white'},
            {'if': {'filter_query': '{color} = Y'},'backgroundColor': 'yellow','color': 'black'},
            {'if': {'filter_query': '{color} = M'},'backgroundColor': 'magenta','color': 'white'},
            {'if': {'filter_query': '{color} = O'},'backgroundColor': 'orange','color': 'white'},
            {'if': {'filter_query': '{color} = C'},'backgroundColor': 'cyan','color': 'white'},
            {'if': {'filter_query': '{color} = B'},'backgroundColor': 'blue','color': 'white'},
            {'if': {'filter_query': '{color} = Z'},'backgroundColor': 'black','color': 'white'},
            {'if': {'filter_query': '{color} = R'},'backgroundColor': 'red','color': 'white'}
        ]

        data = dash_table.DataTable(
            id='table',
            columns=[{"name": i, "id": i} for i in columns],
            data=records,
            filter_action='native',
            sort_action="native",
            page_action='native',
            page_size=30,
            fixed_columns={'headers': True, 'data': nr_of_fixed_columns},
            style_table={'overflowX': 'auto','minWidth':'100%'},
            style_cell_conditional=cell_style,
            style_data_conditional=data_style
        )

        if debug:
            done = round(time.time() * 1000)
            print('It took %d miliseconds to generate this data.' %(done - now))

        return fig,data

    # start server on localhost
    app.run_server(host=ip, port=port, debug=True)

def main(argv=sys.argv):
    __version__ = "0.1"

    parser = argparse.ArgumentParser(
        description="analyzer.py v%s - Fault Injection Analyzer" % __version__,
        prog="analyzer"
    ) 
    parser.add_argument("--directory", help="Database directory", required=True)
    parser.add_argument("--port", help="Server port", required=False, default=8080)
    parser.add_argument("--ip", help="Server address", required=False, default='127.0.0.1')
    parser.add_argument("-x", help="parameter to plot on the x-axis", required=False, default='delay')
    parser.add_argument("-y", help="parameter to plot on the y-axis", required=False, default='length')
    parser.add_argument("--aspect-ratio", help="aspect ratio of the plot relative to x-axis", required=False, default=0, type=float)
    parser.add_argument("--auto-update", help="Whether to update the plot automatically. Optionally pass the update interval in seconds.", required=False, type=int, nargs='?', const=1, default=0)

    args = parser.parse_args()
    run(directory=args.directory, ip=args.ip, port=args.port, x_axis=args.x, y_axis=args.y, aspect_ratio=args.aspect_ratio, auto_update_interval=args.auto_update, debug=True)

if __name__ == "__main__":
    main()
