# coding:utf-8
# @Time     : 2024/8/27 下午3:08
# @Author   : ouyangyouzhang
# @FileName : export_like_msg.py
# @Describe : 导出点赞的消息

import psycopg2
from psycopg2 import sql
import os
import pandas as pd


class ExportMsg:
    def __init__(self):
        try:
            # 连接数据库
            self.connection = psycopg2.connect(
                user="postgres",
                password="difyai123456",
                host="172.20.0.145",
                port=5432,
                database="dify"
            )

        except (Exception, psycopg2.Error) as error:
            print("Error while connecting to PostgreSQL", error)

    def export_like_msg_id(self, user_id: str) -> list:
        cursor = self.connection.cursor()

        # 构建SQL查询
        query = sql.SQL("""
            SELECT message_id, from_end_user_id, rating
            FROM message_feedbacks
            WHERE from_end_user_id = %s AND rating = %s
        """)

        # 执行查询并获取结果
        cursor.execute(query, (user_id, "like"))
        records = cursor.fetchall()
        if cursor:
            cursor.close()

        return [item[0] for item in records]

    def export_QA_to_excel(self, msg_id_list, save_file) -> bool:
        q_list = []
        a_list = []

        cursor = self.connection.cursor()
        for item in msg_id_list:
            # 构建SQL查询
            query = sql.SQL("""
                SELECT id, query, answer
                FROM messages
                WHERE id = %s
            """)

            # 执行查询并获取结果
            cursor.execute(query, (item,))
            records = cursor.fetchall()
            q_list.append(records[0][1])
            a_list.append(records[0][2])
        if cursor:
            cursor.close()

        data = {"query": q_list, "answer": a_list}
        df = pd.DataFrame(data)
        df.to_excel(save_file, index=False)


def main(arg1: list) -> dict:
    if arg1 is None or len(arg1) == 0:
        return {
            "result": [],
        }
    new_list = [{"title": item["title"], "content": item["content"]} for item in arg1]
    return {
        "result": new_list,
    }


if __name__ == '__main__':
    print(e)
    ex = ExportMsg()
    msg_list = ex.export_like_msg_id("de014752-5509-4c9e-a26d-72d7d955f895")
    ex.export_QA_to_excel(msg_list, "D:\\Code\\PGSql\\like_msg.xlsx")
    print('Python')