# TODO: sync pingchat log to serverless every 1 hour. run as a cron job. share chatbot-data/logs
import os
import time
import hashlib
from copy import deepcopy
import argparse
import tarfile
import tempfile
import json
import uuid
from environs import Env
from sqlalchemy import create_engine, text
import pandas as pd

env = Env()
env.read_env()
STORAGE_DSN = env.str("STORAGE_DSN")
store = create_engine(STORAGE_DSN, connect_args={'ssl': {'ssl_ca': '/etc/ssl/certs/ca-certificates.crt'}})

class LogParser:
    def __init__(self, input_path, start_ts=None, end_ts=None):
        self.input_path = input_path
        self.start_ts = start_ts
        self.end_ts = end_ts
        self.temp_dir = None
        self.unexpected_logs = []

    def parse_log_file(self, filepath):
        """Parses a JSON file and returns its contents as a dictionary."""
        print("Parsing file: {}".format(filepath))
        with open(filepath, 'r') as f:
            try:
                logs = [json.loads(line) for line in f \
                        if json.loads(line)['timestamp'] >= self.start_ts \
                            and json.loads(line)['timestamp'] <= self.end_ts]
            except:
                print(filepath)
                raise ValueError("Invalid JSON file")

        for i, log in enumerate(logs):
            try:
                msg = json.loads(log['message'])
                logs[i]['message'] = msg
            except:
                raise ValueError("Invalid JSON file")
        return logs

    def parse_log_files(self, filepaths):
        """Parses a list of JSON filepaths and returns their contents as a list of dictionaries."""
        logs = []
        for filepath in filepaths:
            if 'last.json' in filepath:
                continue
            sub_logs = self.parse_log_file(filepath)
            for log in sub_logs:
                logs.append(log)
        return logs

    def get_files_in_directory(self, directory):
        """Returns a list of JSON filepaths in a directory with time limit."""
        files = []
        for file in os.listdir(directory):
            if os.path.isfile(os.path.join(directory, file)) and file.startswith('chatbot'):
                if os.path.getmtime(os.path.join(directory, file)) < self.start_ts:
                    continue
                files.append(file)
        return [os.path.join(directory, f) for f in files]

    def get_files_in_tar_archive(self, archive_path):
        """Returns a list of JSON filepaths in a tar archive without time limit."""
        json_files = []
        with tarfile.open(archive_path, 'r:gz') as tar:
            self.temp_dir = tempfile.TemporaryDirectory()
            tar.extractall(path=self.temp_dir.name)
            for member in tar.getmembers():
                if member.isfile():
                    json_files.append(self.temp_dir.name + '/' + member.name)
        return json_files

    def parse(self):
        if os.path.isdir(self.input_path):
            files = self.get_files_in_directory(self.input_path)
        elif os.path.isfile(self.input_path) and self.input_path.endswith('.tar.gz'):
            files = self.get_files_in_tar_archive(self.input_path)
        elif os.path.isfile(self.input_path):
            files = [self.input_path]
        else:
            raise TypeError("Invalid input format, must be a path to a directory, a compressed tar archive, or a single JSON-like log file")  # noqa

        logs = self.parse_log_files(files)
        if self.temp_dir is not None:
            self.temp_dir.cleanup()
        self.logs_sorted = sorted(logs, key=lambda x: x['timestamp'])
        return self.logs_sorted

    def group_by(self, key):
        def all_elements_in_candidate_list(input_list, candidate_list):
            for element in input_list:
                if element not in candidate_list:
                    return False
            return True
        candidates = ['normal', 'dislike', 'like', 'error']
        if not all_elements_in_candidate_list(key, candidates):
            raise NotImplementedError('Only normal, dislike, like, and error are supported. Input must be a list of these strings, such as ["dislike"] or ["error", "dislike"]')  # noqa
        results = {k: [] for k in key}
        for log in self.logs_sorted:
            for k in key:
                if k == 'normal' and isinstance(log['message'], dict) and log['message'].get('request') is not None:
                    results[k].append(log)
                elif k == 'error' and isinstance(log['message'], str):
                    results[k].append(log)
                elif (k == 'like' or k == 'dislike') and isinstance(log['message'], dict) and log['message'].get('feedback') == k:
                    results[k].append(log)
        return results

def preprocess(classified: dict):
    for k, v in classified.items():
        if k == 'normal':
            for dialog in v:
                tmp = deepcopy(dialog['message']['request'])
                tmp.append({'role': 'assistant', 'content': dialog['message']['response']['content']})
                dialog['messages'] = tmp
                dialog['user'] = dialog['message'].get('user')
                dialog['source'] = dialog['message'].get('source')
                dialog['request'] = dialog['message']['request']
                dialog['response'] = dialog['message']['response']['content']
                dialog['recalls'] = dialog['message']['response']['recalls']
                dialog['revised_q'] = dialog['message']['response']['instruction']
                dialog['links'] = dialog['message']['response']['links']
                dialog['trace_id'] = dialog.get('x-trace-id')
                dialog['qa_hash'] = hashlib.md5(json.dumps(dialog['messages']).encode()).hexdigest() # normal has a slide window with past 2 qa and current qa
                del dialog['message']
        elif k == 'dislike':
            for piece in v:
                piece['user'] = piece['message'].get('user')
                piece['messages'] = piece['message']['messages']['messages']
                piece['description'] = piece['message']['messages'].get('description')
                piece['type'] = piece['message']['messages'].get('type')
                piece['correct_links'] = piece['message']['messages'].get('related_link')
                piece['trace_id'] = piece.get('x-trace-id')
                piece['qa_hash'] = hashlib.md5(json.dumps(piece['messages'][-6:]).encode('utf-8')).hexdigest() # dislike has no window
                del piece['message']
        elif k == 'like':
            for piece in v:
                piece['user'] = piece['message'].get('user')
                piece['messages'] = piece['message']['messages']
                piece['trace_id'] = piece.get('x-trace-id')
                piece['qa_hash'] = hashlib.md5(json.dumps(piece['messages'][-6:]).encode('utf-8')).hexdigest() # like has no window
                del piece['message']
        elif k == 'error':
            for piece in v:
                piece['trace_id'] = piece.get('x-trace-id')

def convert_dict_columns_to_json(df):
    for column in df.columns:
        if df[column].dtype == 'object' and any(isinstance(x, (dict, list)) for x in df[column] if x is not None):
            df[column] = df[column].apply(json.dumps)
    return df

def save(df: pd.DataFrame, remote: bool, dest: str):
    if remote:
        df = convert_dict_columns_to_json(df)
        df.to_sql(dest, con=store, if_exists='append', index=False)
            
    else:
        # check dest.csv exists
        if os.path.exists(dest+'.csv'):
            raise FileExistsError('File already exists')
        else:
            df.to_csv(dest+'.csv', index=False)
            print('Successfully saved to {}'.format(dest+'.csv'))  # TODO: change to logging

def select_latest_row(group):
    if len(group) == 1:
        return group
    non_empty = group[group['type'].notna() & (group['type'] != '')]
    if not non_empty.empty:
        return non_empty.nlargest(1, 'created_at')
    return group.nlargest(1, 'created_at')

def delete_older_ids(df_a, df_b, condition):
    """delete the older condition rows that both exist in df_a and df_b. the condition in a or b is unique."""
    # condition is a list of column names
    merged_df = pd.merge(df_a, df_b, on=condition, how='outer', suffixes=('_a', '_b'))
    # get both exist rows
    common_ids = merged_df.dropna(subset=['created_at_a', 'created_at_b'])

    for _, row in common_ids.iterrows():
        if row['created_at_a'] > row['created_at_b']:
            # delete from df_b
            mask = True
            for col in condition:
                mask = mask & (df_b[col] != row[col])
            df_b = df_b[mask]
        else:
            # delete from df_a
            mask = True
            for col in condition:
                mask = mask & (df_a[col] != row[col])
            df_a = df_a[mask]
    return df_a, df_b

def update_status(df, status):
    conn = store.connect()
    for _, row in df.iterrows():
        trace_id = row.get('trace_id')
        created_at = row.get('created_at')
        _id = row.get('id')
        query = f"""UPDATE chatbot_opt.qa SET `status` = '{status}', `link_id` = '{_id}' WHERE `id` IN (SELECT `id` FROM chatbot_opt.qa qa WHERE qa.trace_id = '{trace_id}' and qa.created_at < '{created_at}' ORDER BY qa.created_at DESC LIMIT 1);"""  # slack has duplicate trace_id
        conn.execute(text(query))
        # # used for old pingchat that has no trace_id before Jun 29, 2023
        # qa_hash = row.get('qa_hash')
        # user = row.get('user')
        # query_for_pingchat = f"""UPDATE TABLE chatbot_opt.qa SET `status` = '{status}' WHERE `id` IN (SELECT `id` FROM chatbot_opt.qa qa WHERE qa.qa_hash = '{qa_hash}' and qa.`user` = '{user}' and qa.created_at < '{created_at}') ORDER BY qa.created_at ASC LIMIT 1;"""
        # conn.execute(text(query_for_pingchat))
        conn.commit()
    conn.close()

def clear_duplicates(like_df, dislike_df):
    if not like_df.empty:
        like_df.sort_values(by='created_at', ascending=False, inplace=True)
        like_df_w_id = like_df[like_df['trace_id'].notna()]
        like_df_w_id = like_df_w_id.drop_duplicates(subset=['trace_id'], keep='first')
        # like_df_w_id = like_df_w_id.drop_duplicates(subset=['qa_hash', 'user'], keep='first') # used for pingchat.pingcap.net that has random trace_id
        # # used for old pingchat that has no trace_id before Jun 29, 2023
        # like_df_wo_id = like_df[like_df['trace_id'].isna()]
        # like_df_wo_id = like_df_wo_id.drop_duplicates(subset=['qa_hash', 'user'], keep='first')

    if not dislike_df.empty:
        dislike_df.sort_values(by='created_at', ascending=False, inplace=True)
        dislike_df_w_id = dislike_df[dislike_df['trace_id'].notna()]
        dislike_df_w_id = dislike_df_w_id.groupby('trace_id').apply(select_latest_row).reset_index(drop=True)
        # dislike_df_w_id = dislike_df_w_id.groupby(['qa_hash', 'user']).apply(select_latest_row).reset_index(drop=True) # used for pingchat.pingcap.net that has random trace_id
        # # used for old pingchat that has no trace_id before Jun 29, 2023
        # dislike_df_wo_id = dislike_df[dislike_df['trace_id'].isna()]
        # dislike_df_wo_id = dislike_df_wo_id.groupby(['qa_hash', 'user']).apply(select_latest_row).reset_index(drop=True)

    if not like_df.empty and not dislike_df.empty:
        like_df_w_id, dislike_df_w_id = delete_older_ids(like_df_w_id, dislike_df_w_id, ['trace_id'])
        # like_df_w_id, dislike_df_w_id = delete_older_ids(like_df_w_id, dislike_df_w_id, ['qa_hash', 'user']) # used for pingchat.pingcap.net that has random trace_id
        # # used for old pingchat that has no trace_id before Jun 29, 2023
        # like_df_wo_id, dislike_df_wo_id = delete_older_ids(like_df_wo_id, dislike_df_wo_id, ['qa_hash', 'user'])

    if not like_df.empty:
        # user for new pingchat that has trace_id after Jun 29, 2023
        new_like_df = like_df_w_id
        # # used for old pingchat that has no trace_id before Jun 29, 2023
        # new_like_df = pd.concat([like_df_w_id, like_df_wo_id], axis=0, ignore_index=True)
        # new_like_df.sort_values(by='created_at', ascending=False, inplace=True)
    else:
        new_like_df = pd.DataFrame()
    if not dislike_df.empty:
        # user for new pingchat that has trace_id after Jun 29, 2023
        new_dislike_df = dislike_df_w_id
        # # used for old pingchat that has no trace_id before Jun 29, 2023
        # new_dislike_df = pd.concat([dislike_df_w_id, dislike_df_wo_id], axis=0, ignore_index=True)
        # new_dislike_df.sort_values(by='created_at', ascending=False, inplace=True)
    else:
        new_dislike_df = pd.DataFrame()
    return new_like_df, new_dislike_df


if __name__ == '__main__':
    # python sync_log.py chatbot.json
    # TODO: time range with cron job
    parser = argparse.ArgumentParser(description='Parse JSON files')
    parser.add_argument('input', metavar='INPUT', type=str,
                        help='a path to a directory, a compressed tar archive, or a single JSON file')
    args = parser.parse_args()

    last_ts_file = '/data/logs/last.json'
    if not os.path.exists(last_ts_file):
        with open(last_ts_file, 'w') as f:
            json.dump({'last_ts': 1682632800.0}, f)  # the logs before the time has NaN bugs
    # load the last timestamp from the file
    with open(last_ts_file, 'r') as f:
        data = json.load(f)
        start_ts = data['last_ts']
    # get the current timestamp
    end_ts = time.time()
    
    parser = LogParser(args.input, start_ts, end_ts)
    logs = parser.parse()

    sub = parser.group_by(['dislike', 'like', 'normal', 'error'])

    preprocess(sub)

    diss = sub['dislike']
    likes = sub['like']
    dialogs = sub['normal']
    errors = sub['error']

    dislike_df = pd.DataFrame.from_records(diss)
    if len(dislike_df) > 0:
        dislike_df['id'] = dislike_df.apply(lambda x: uuid.uuid4(), axis=1)
        dislike_df['created_at'] = pd.to_datetime(dislike_df['timestamp'], unit='s')
        dislike_df.drop(columns=['level', 'traceback', 'timestamp', 'x-trace-id'], inplace=True)

    normal_df = pd.DataFrame.from_records(dialogs)
    if len(normal_df) > 0:
        normal_df['id'] = normal_df.apply(lambda x: uuid.uuid4(), axis=1)
        normal_df['created_at'] = pd.to_datetime(normal_df['timestamp'], unit='s')
        normal_df['status'] = normal_df['traceback'].apply(lambda x: 'error' if x is not None else 'normal')
        normal_df['link_id'] = None
        normal_df.drop(columns=['level', 'messages', 'timestamp', 'x-trace-id'], inplace=True)

    like_df = pd.DataFrame.from_records(likes)
    if len(like_df) > 0:
        like_df['id'] = like_df.apply(lambda x: uuid.uuid4(), axis=1)
        like_df['created_at'] = pd.to_datetime(like_df['timestamp'], unit='s')
        like_df.drop(columns=['level', 'traceback', 'timestamp', 'x-trace-id'], inplace=True)

    error_df = pd.DataFrame.from_records(errors)
    if len(error_df) > 0:
        error_df['id'] = error_df.apply(lambda x: uuid.uuid4(), axis=1)
        error_df['created_at'] = pd.to_datetime(error_df['timestamp'], unit='s')
        error_df.drop(columns=['level', 'timestamp', 'x-trace-id'], inplace=True)
        # error_df.drop(columns=['level', 'timestamp'], inplace=True) # before Jun 29, 2023

    save(dislike_df, remote=True, dest='dislike')
    save(normal_df, remote=True, dest='qa')
    save(like_df, remote=True, dest='like')
    save(error_df, remote=True, dest='error')

    like_df, dislike_df = clear_duplicates(like_df, dislike_df)

    update_status(dislike_df, 'dislike')
    update_status(like_df, 'like')

    data['last_ts'] = end_ts
    with open(last_ts_file, 'w') as f:
        json.dump(data, f)