import os
import time
import datetime
import json
import shutil
import pandas as pd
import logging

from pathlib import Path
from tqdm import tqdm

from src.My_Logger.project_logger import LoggerClass
from src.PostgresSQL.databaseHandler import DatabaseHandler
from src.utils.utils import load_config, size_of_file

project_path = Path(__file__).parent.parent.parent.absolute()
data_path = os.path.join(project_path, 'data')
config_path = os.path.join(project_path, 'config')


def process_message_text(text):
    if isinstance(text, str):
        if '[Sticker ' in text and ']' in text:
            return 'Sticker'
        return text
    elif isinstance(text, list):
        res = ''
        for item in text:
            if isinstance(item, str):
                res += item + '\n'
            elif isinstance(item, dict):
                res += f"{item.get('text')}\n"
            elif not item:
                continue
            else:
                raise TypeError
        return res
    else:
        raise TypeError


class MessageUploader:
    def __init__(self, config_path, local_logger, local_load_path, local_file_db_path, size_limit=0.5):
        self.logger = local_logger
        self.config = load_config(config_path)
        self.load_path = local_load_path
        self.file_param_list = ['photo', 'file']
        self.db_handler = DatabaseHandler(self.config, self.logger)
        self.file_db_path = local_file_db_path

        self.size_limit = size_limit * 1024 * 1024 * 1024
        self.single_file_size_limit = 0.24 * 1024 * 1024 * 1024  # 240MB

        self.table_format = {
            'message_id': 'text',
            'message_type': 'text',
            'message_datetime': 'text',
            'message_unix': 'text',
            'message_from': 'text',
            'message_from_id': 'text',
            'message_text': 'text',
            'forward_from': 'text',
            'reply_message_id': 'text',
            'file_location': 'text',
            'file_type': 'text',
            'file_size': 'text',
            'file_size_raw': 'numeric',
        }

    def do_upload(self):
        start_time = time.time()
        loaded_messages = self.load_message()
        table_name = f'{loaded_messages["name"]}_{loaded_messages["type"]}'
        if 'rynco.me' in table_name:
            table_name = 'Rynco_Group_private_supergroup'
        self.prepare_database(table_name)
        target_list = self.process_loaded_message(loaded_messages, table_name)
        self.upload_to_database(target_list, table_name)
        self.logger.info(f'Message upload completed in {round(time.time() - start_time, 3)}s')

    def prepare_database(self, table_name):
        # self.db_handler.create_database(table_name, force_create=False)
        if not self.db_handler.exists_table(table_name):
            self.db_handler.create_table(self.table_format, table_name, False, ['message_id'])
            time.sleep(0.5)
            self.logger.info(f'Table {table_name} created.')

    def load_message(self):
        with open(os.path.join(self.load_path, 'result.json'), 'r', encoding='utf-8') as file:
            file_content = file.read()
        self.logger.info('Message loaded')
        return json.loads(file_content)

    def upload_to_database(self, target_list, table_name):
        db_handler = self.db_handler
        db_handler.insert_json_into_table(table_name=table_name, data_list=target_list,
                                          batch=10000, unique_columns='message_id')

    def process_loaded_message(self, loaded_messages, table_name):
        self.logger.info('Start processing messages')
        start_time = time.time()
        conversation_id = loaded_messages['id']
        messages = [item for item in loaded_messages['messages']]
        target_list = []
        for message in tqdm(messages, desc=f'Processing messages'):
            if message['type'] == 'service':
                continue
            target = {
                'message_id': message['id'],
                'message_type': message['type'],
                'message_datetime': message['date'],
                'message_unix': message['date_unixtime'],
                'message_from': message['from'],
                'message_from_id': message['from_id'],
                'message_text': process_message_text(message['text']),
                'forward_from': message.get('forwarded_from', None),
                'reply_message_id': message.get('reply_to_message_id', None),
                'file_location': None,
                'file_type': None,
                'file_size': None,
                'file_size_raw': None
            }
            for file_param in self.file_param_list:
                if file_param in message:
                    file_location = self.transfer_file_to_db(message[file_param])
                    target['file_location'] = file_location
                    target['file_size_raw'] = os.path.getsize(os.path.join(self.load_path, message[file_param]))
                    target['file_size'] = size_of_file(target['file_size_raw'])
                    if file_param == 'photo':
                        target['file_type'] = file_param
                    elif 'sticker_emoji' in message:
                        target['file_type'] = 'sticker'
                    elif 'media_type' in message:
                        target['file_type'] = message['media_type']
                    else:
                        target['file_type'] = 'file'

            target_list.append(target)
        self.logger.info(f'Data process complete, time: {time.time() - start_time}')
        return target_list

    def transfer_file_to_db(self, file_path):
        source_path = os.path.join(self.load_path, file_path)
        destination_path = os.path.join(self.file_db_path, file_path)
        destination_folder, file_name = os.path.split(destination_path)
        new_file_name = file_name
        # 路径不存在则创建路径
        if not os.path.exists(destination_folder):
            os.mkdir(destination_folder)

        # 已存在相同文件名的情况
        if os.path.exists(destination_path):
            # 如果文件大小不相等，或者大小相同但不是sticker文件，直接进行复制操作。
            if (os.path.getsize(source_path) != os.path.getsize(destination_path) or
                    ('sticker' not in source_path) and
                    os.path.getsize(source_path) == os.path.getsize(destination_path)):
                i = 0
                while os.path.exists(destination_path):
                    new_file_name = f"{i}_{file_name}"
                    destination_path = os.path.join(destination_folder, new_file_name)
                    i += 1
            # 如果是sticker文件，则直接进行映射（sticker文件出错问题应该不大）

        shutil.copy2(source_path, destination_path)
        return file_path.replace(file_name, new_file_name)

    def logger_test(self):
        self.logger.info('this is a test')


def upload_didi_saved_message_to_db():
    msg_list = []
    csv_content = pd.read_csv('../data/Rynco_Group.csv', encoding='utf-8')
    for row, content in csv_content.iterrows():
        try:
            text = process_message_text(content['text'])
        except TypeError as e:
            text = ''
        target = {
            'message_id': content['id'],
            'message_type': '' + 'message',
            'message_datetime': content['date'],
            'message_unix': int(datetime.datetime.strptime(content['date'], '%Y-%m-%dT%H:%M:%S').timestamp()),
            'message_from': content['from_name'],
            'message_from_id': content['from_id'],
            'message_text': text,
            'forward_from': content.get('forwarded_from'),
            'reply_message_id': content.get('reply_to_message_id', None),
            'file_location': None,
            'file_type': 'sticker' if 'Sticker' == text else None,
            'file_size': None,
            'file_size_raw': None
        }
        msg_list.append(target)
    tg.upload_to_database(msg_list, 'Rynco_Group_private_supergroup')


if __name__ == "__main__":

    config = load_config(config_path)
    config['database'] = "TestDB"
    tg = MessageUploader(config_path=config_path,
                         local_logger=LoggerClass().get_log(logging.DEBUG),
                         local_load_path='E:\\SNSFiles\\ChatExport_2024-07-10',
                         local_file_db_path=config['file_db_path'],
                         size_limit=0.25)
    try:
        tg.do_upload()
        # tg.logger_test()
    finally:
        tg.db_handler.close()
        print('done')

