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.db_handler import PgHandler
from src.utils.utils import size_of_file
from src.config.load_config import project_config

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, local_logger, local_load_path, local_file_db_path, size_limit=0.5):
        self.logger = local_logger
        self.config = project_config
        self.load_path = local_load_path
        self.file_param_list = ['photo', 'file']
        # 使用 PgHandler 替代旧的 DatabaseHandler，按 config 构造连接
        self.db_handler = PgHandler(host=self.config['host'], port=int(self.config.get('port', 5432)),
                                    user=self.config.get('user'), password=self.config.get('password'),
                                    database=self.config.get('database'))
        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):
            # 使用 PgHandler 的 create_table_from_dict 接口创建表
            self.db_handler.create_table_from_dict(self.table_format, table_name, force_create=False, primary_key_list=['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
        # 调用 PgHandler.insert_json_into_table(table, json_list, schema, batch, unique_columns)
        db_handler.insert_json_into_table(table=table_name, json_list=target_list, schema='public', 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:
                    # 尝试复制文件到文件库，transfer_file_to_db 在失败时返回 None
                    file_path = message[file_param]
                    file_location = self.transfer_file_to_db(file_path)
                    target['file_location'] = file_location
                    if file_location:
                        source_path = os.path.join(self.load_path, file_path)
                        try:
                            file_size_raw = os.path.getsize(source_path)
                            target['file_size_raw'] = file_size_raw
                            target['file_size'] = size_of_file(file_size_raw)
                        except Exception:
                            # 无法获取大小，置为 None 并记录
                            self.logger.warning(f'Could not get size for source file: {source_path}')
                            target['file_size_raw'] = None
                            target['file_size'] = None
                        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'
                    else:
                        # 复制失败或文件不存在，保持 file_* 字段为 None
                        self.logger.warning(f'File not copied, skipped file field for message id {message.get("id")}')

            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):
        # 防御性检查：file_path 不是字符串或是占位文本时直接返回 None
        try:
            if not isinstance(file_path, str):
                self.logger.warning(f'Invalid file path type: {type(file_path)}')
                return None
            low = file_path.lower()
            if 'exceeds maximum size' in low or 'file exceeds' in low:
                self.logger.warning(f'File path appears to be placeholder (too large): {file_path}')
                return None

            source_path = os.path.join(self.load_path, file_path)
            if not os.path.exists(source_path):
                self.logger.warning(f'Source file does not exist: {source_path}')
                return None

            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

            # 确保目标目录存在
            os.makedirs(destination_folder, exist_ok=True)

            # 已存在相同文件名的情况
            if os.path.exists(destination_path):
                try:
                    src_size = os.path.getsize(source_path)
                    dst_size = os.path.getsize(destination_path)
                except Exception:
                    src_size = dst_size = None
                # 如果文件大小不相等，或大小相同但非 sticker，改名复制
                if (src_size is not None and dst_size is not None and src_size != dst_size) or (
                        src_size is not None and dst_size is not None and src_size == dst_size and 'sticker' not in source_path):
                    i = 0
                    base_name = file_name
                    while os.path.exists(destination_path):
                        new_file_name = f"{i}_{base_name}"
                        destination_path = os.path.join(destination_folder, new_file_name)
                        i += 1

            try:
                shutil.copy2(source_path, destination_path)
            except Exception as e:
                self.logger.warning(f'Failed to copy file {source_path} to {destination_path}: {e}')
                return None

            return file_path.replace(file_name, new_file_name)
        except Exception as e:
            self.logger.exception(f'Unexpected error in transfer_file_to_db for {file_path}: {e}')
            return None

    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__":
    tg = MessageUploader(
        local_logger=LoggerClass().get_log(logging.DEBUG),
        local_load_path='E:\\SNSFiles\\ChatExport_2026-01-01',
        local_file_db_path=project_config['file_db_path'],
        size_limit=0.25
    )
    try:
        tg.do_upload()
        # tg.logger_test()
    finally:
        tg.db_handler.close()
        print('done')
