"""
筛选  整合  线上的   口语 核心互动   对话记录数据

数据筛选流程如下:
一 步骤一
首先， 在 PGsql数据库中  筛选出  口语核心互动对应的  session_id.
数据库相关配置  从.env中读取:
PG_DB_HOST = xxx
PG_DB_PORT = xxx
PG_DB_USER = xxx
PG_DB_PASSWORD = xxx
PG_DB_DATABASE = xxx

读取以下数据表: 
user_component_play_record_0 ~ user_component_play_record_7

支持输入时间范围
起始时间 和 截止时间   配置格式:   "20250110"

数据表中的时间字段为 updated_at , 格式样例: "2025-11-05 19:35:46.698246+08:00"

在这些时间范围内，筛选以下数据:
c_type 为 core_speaking_reply 或者 core_speaking_inquiry 的数据

输出总的数据条数

然后导出 中间 excel文件

包含以下字段：
user_id,
session_id,
c_type,
c_id,
play_result,
updated_at

二. 步骤二
根据 c_type 和 c_id 筛选核心互动的配置  补充一些字段。

需要读取配置表:
mysql表  core_interaction_component
相关环境变量在.env:
MYSQL_HOST=xxx
MYSQL_USERNAME=xxx
MYSQL_PASSWORD=xxx
MYSQL_DATABASE=xxx
MYSQL_PORT=xxx

基于 c_type 和 c_id 字段匹配， 在 步骤一表格内容基础上追加以下字段:
title
reference_dialog 从 component_config 中抽取出 reference_dialog 字段的内容。
component_config内容样例:
```
{"taskInfo":{"cId":"0000001","cType":"core_speaking_inquiry","title":"询问种植甜瓜的信息","taskDesc":"向Ben提问甜瓜种植的最佳季节、浇水频率和成熟的季节；","sceneDesc":"我和Ben到甜味城，参观了水果资源站和种植园。Ben的妈妈Kate讲了种植知识，我们都很感兴趣，想一起种甜瓜。我不懂，便问Ben，他虽没种过、不确定，还是告诉我注意事项。","img":"","key":[{"desc":"询问种植信息","keyList":[{"type":"default","npcId":269,"content":"Have you ever planted a ...?","desc":"你种过......吗？"},{"type":"default","npcId":269,"content":"What season is the best time to plant ...?","desc":"种植......的最佳时间是哪个季节？"},{"type":"default","npcId":269,"content":"Do ... need ... every day?","desc":"......需要每天浇......么？"}]}]},"dialogSetting":{"setting":{"npcName":"Ben","npcId":287,"round":5,"checkRound":3}},"dialogConfig":{"config":{"asrPrompt":"melon,summer,autumn,water,frequency,plant,season,harvest","promptInfo":{"default":"# 1. 角色（你要扮演谁）\n- 你是 Ben，一个 8 岁的小男孩，对种植水果感兴趣但不太确定具体细节。\n- 语言风格：简单、直接，偶尔带有不确定的语气。\n- 示例表达：\n  - \"I think summer. It's warm then.\"\n  - \"Maybe every two days? Not every day, I think.\"\n\n# 2. 任务（你如何参与到整个对话）\n- 你需要只在用户提问时提供信息，不会主动提及种植甜瓜的具体细节。\n- 如果用户提问相关内容，你需要根据知识库中的信息回答，不编造或偏离。\n- 如果用户的问题不清晰，你需要尝试澄清后再作答。\n- 如果用户长时间不提问或偏离主题，你需要温和、自然地进行交谈，引导回到主题。\n- 当所有知识点已传达后，你需要鼓励用户开始行动。\n\n# 3. 背景信息（引用配置）\n`你是 Ben，你和用户来到了甜味城Sweet Town。你们参观了水果资源站和种植园。在种植园中，你的妈妈Kate给你和用户介绍了一些种植水果的知识。你和用户对此很感兴趣。你们想要一起种一颗甜瓜。用户不知道种植甜瓜的知识，于是向你提问。虽然你没有种过甜瓜，对什么都不确定。但你还是回答了用户的问题，告诉用户关于种植甜瓜需要注意的事情。`\n\n# 4. 知识库（你知道的信息）\n- 种甜瓜的最佳季节：应该在夏天\n- 种甜瓜的浇水频率：应该隔一天浇一次水\n- 甜瓜成熟的季节：秋天\n- 如果用户提问相关内容，你会用这些信息来回答。\n\n# 5. 语言风格（固定内容）\n  1.  使用标准、正式的英语，水平为 CEFR A1/A2，每句话不超过 10 个单词\n  2.  始终保持礼貌和友好\n  3.  尽量避免重复表达，适当变换措辞\n\n# 6. 开场白\n你由你开始对话，你会说：“Let's plant a melon now! Or do you still have some questions？”\n\n# 7. 回应方式（固定内容）\n`你只在用户提问时才根据知识库中的信息回答。其他时候，以符合你身份的方式，自然地进行交谈。不主动提供信息，不偏离语境。`","final_goal":"Ben 说出了种植种甜瓜的最佳季节应该在夏天、种甜瓜的浇水频率应该隔一天浇一次水以及甜瓜成熟的季节是秋天","in_progress_goal":"Ben 说出了种植种甜瓜的最佳季节应该在夏天、种甜瓜的浇水频率应该隔一天浇一次水以及甜瓜成熟的季节是秋天","reference_dialog":"# 示例对话\nBen: Let's plant a melon now! Or do you still have some questions？\nYou: Great! When is the best time to plant it?\nBen: I think summer. It's warm then. Mom said melons like warm weather.\nYou: Oh, good. How often should we water it?\nBen: Maybe every two days? Not every day, I think.\nYou: And when will it be ready to eat?\nBen: Autumn, I guess. Plant in summer, get melons in autumn. That sounds right.","scene":"#任务背景\n你是 Ben，你和用户来到了甜味城Sweet Town。你们参观了水果资源站和种植园。在种植园中，你的妈妈Kate给你和用户介绍了一些种植水果的知识。你和用户对此很感兴趣。你们想要一起种一颗甜瓜。用户不知道种植甜瓜的知识，于是向你提问。虽然你没有种过甜瓜，对什么都不确定。但你还是回答了用户的问题，告诉用户关于种植甜瓜需要注意的事情。","user_knowledge":"# 知识\n- 询问种植信息\nHave you ever planted a ...?  你种过......吗？\nWhat season is the best time to plant ...?  种植......的最佳时间是哪个季节？\nDo ... need ... every day?  ......需要每天浇......么？","user_scene":"我和Ben到甜味城，参观了水果资源站和种植园。Ben的妈妈Kate讲了种植知识，我们都很感兴趣，想一起种甜瓜。我不懂，便问Ben，他虽没种过、不确定，还是告诉我注意事项。","user_task":"向Ben提问甜瓜种植的最佳季节、浇水频率和成熟的季节；"}}},"studyInfo":{"learningPart":{"learning":[{"question":{"desc":"现在你需要询问Ben关于种植甜瓜的最佳季节。"},"optionList":[{"option":"When is the best time to plant it?","feedbackDesc":"太棒了！你正确地询问了种植甜瓜的最佳季节。请大声朗读这句话！"},{"option":"How often should we water it?","feedbackDesc":"这句话是询问浇水频率的，不是询问最佳种植季节的。请再试一次，询问Ben种植甜瓜的最佳季节。"},{"option":"When will it be ready to eat?","feedbackDesc":"这句话是询问甜瓜成熟季节的，不是询问最佳种植季节的。请再试一次，询问Ben种植甜瓜的最佳季节。"}],"answer":[0],"read":{"type":"user","npcId":30,"content":"When is the best time to plant it?"},"feedback":{"type":"npc","npcName":"Ben","npcId":287,"content":"I think summer. It's warm then. Mom said melons like warm weather."}},{"question":{"desc":"Ben告诉你种植甜瓜的最佳季节是夏天。现在你需要询问Ben关于种植甜瓜的浇水频率。"},"optionList":[{"option":"When is the best time to plant it?","feedbackDesc":"这句话是询问最佳种植季节的，不是询问浇水频率的。请再试一次，询问Ben种植甜瓜的浇水频率。"},{"option":"How often should we water it?","feedbackDesc":"太棒了！你正确地询问了种植甜瓜的浇水频率。请大声朗读这句话！"},{"option":"When will it be ready to eat?","feedbackDesc":"这句话是询问甜瓜成熟季节的，不是询问浇水频率的。请再试一次，询问Ben种植甜瓜的浇水频率。"}],"answer":[1],"read":{"type":"user","npcId":30,"content":"How often should we water it?"},"feedback":{"type":"npc","npcName":"Ben","npcId":287,"content":"Maybe every two days? Not every day, I think."}},{"question":{"desc":"Ben告诉你种植甜瓜的浇水频率是隔一天一次。现在你需要询问Ben关于甜瓜成熟的季节。"},"optionList":[{"option":"When is the best time to plant it?","feedbackDesc":"这句话是询问最佳种植季节的，不是询问甜瓜成熟季节的。请再试一次，询问Ben甜瓜成熟的季节。"},{"option":"How often should we water it?","feedbackDesc":"这句话是询问浇水频率的，不是询问甜瓜成熟季节的。请再试一次，询问Ben甜瓜成熟的季节。"},{"option":"When will it be ready to eat?","feedbackDesc":"太棒了！你正确地询问了甜瓜成熟的季节。请大声朗读这句话！"}],"answer":[2],"read":{"type":"user","npcId":30,"content":"When will it be ready to eat?"},"feedback":{"type":"npc","npcName":"Ben","npcId":287,"content":"Autumn, I guess. Plant in summer, get melons in autumn. That sounds right."}}],"opening":{"type":"npc","npcName":"Ben","npcId":287,"content":"Let's plant a melon now! Or do you still have some questions？","desc":"Ben邀请你一起种植甜瓜，并询问你是否还有问题。"},"closing":{"desc":"Ben已经回答了所有关于种植甜瓜的问题，任务成功完成！"}}},"kpInfoList":[{"kpId":"","kpType":"sentence","kpTitle":"What do you think about the fight?","kpSkill":"sentence_pron","kpSkillName":"语音"},{"kpId":"","kpType":"sentence","kpTitle":"What do you think about the fight?","kpSkill":"sentence_meaning","kpSkillName":"语义"},{"kpId":"","kpType":"sentence","kpTitle":"Can you help us?","kpSkill":"sentence_pron","kpSkillName":"语音"},{"kpId":"","kpType":"sentence","kpTitle":"Can you help us?","kpSkill":"sentence_meaning","kpSkillName":"语义"},{"kpId":"","kpType":"sentence","kpTitle":"Do you know any way to beat him?","kpSkill":"sentence_pron","kpSkillName":"语音"},{"kpId":"","kpType":"sentence","kpTitle":"Do you know any way to beat him?","kpSkill":"sentence_meaning","kpSkillName":"语义"}]}
```

追加后，excel文件包含以下字段:
user_id,
session_id,
c_type,
c_id,
play_result,
updated_at,
title
reference_dialog

三. 步骤三  追加对话历史数据
对话历史数据，需要根据以下es数据库来补充:

es索引: llm_roleplayagent_round_log
相关环境变量在.env:
ES_HOST=xxx
ES_PORT=xxx
ES_SCHEME=xxx
ES_USER=xxx
ES_PASSWORD=xxx

基于每条记录中的 session_id, 匹配 es日志中 session_id 相同 且 action为 get_chat 对应的记录，整理后， 追加为 chat_log 字段。

es中的日志是每轮作为一条记录，按以下逻辑进行拼接:
读取 current_round,
current round 为 0 ， 则 chat_log中加入 npc_message 的内容   "npc: " + npc_message
current round 为 1~n  按顺序 依次追加  user_input 和 npc_message ， 每轮之间用换行符隔开。
完全拼接后  最为 chat_log 内容
完整样例:
```
npc:xxx
user:xxx 
npc:xxx
... ...
```

拼接完成后  追加  chat_log  和  round_num （取最大的current_round）

最终输出的 excel文件字段:
user_id,
session_id,
c_type,
c_id,
play_result,
updated_at,
title,
reference_dialog,
chat_log,
user_behavior_info,
round_num


----------------------
根据以上需求  提供一个数据处理的脚本   尽量用高效的匹配。  我只需要输出最终的匹配文件，一个简单的功能脚本。 脚本不需要太复杂。但在输出的节点增加必要的日志  方便我了解数据量和进度 输入 时间范围 在 脚本开头配置即可。
----------------------

补充需求:
pg sql数据库中 增加字段  user_behavior_info 读取。 并保留到最终的输出excel文档中  在 chat_log字段之后。  其他不变。
----------------------

补充需求 25.11.07:
从 mysql表中 额外读取两个字段的信息进行处理，

1. lesson
抽取related_path字段中的lessonIndex内容 (4):
{"packageId":2,"unitId":26,"lessonId":128,"packageIndex":1,"unitIndex":12,"lessonIndex":4}

2. knowledge_points
直接读取 kp_relation_info 的内容。

3. in_progress_goal
读取 和 reference_dialog 平级的  in_progress_goal 字段内容。

4. final_goal
读取 和 reference_dialog 平级的  final_goal 字段内容。

以上四个字段 都追加到最终输出的表中，
全部输出字段顺序如下:

user_id,
session_id,
c_type,
c_id,
play_result,
updated_at,
title,
lesson,
knowledge_points,
in_progress_goal,
final_goal,
reference_dialog,
chat_log,
user_behavior_info,
round_num


"""

import os
import json
import pandas as pd
import psycopg2
import pymysql
from elasticsearch import Elasticsearch
from datetime import datetime
from dotenv import load_dotenv
import logging

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# 时间范围配置 - 修改这里的日期范围
START_DATE = "20251001"  # 起始时间 格式: "20250110"
END_DATE = "20251031"    # 截止时间 格式: "20250131"

class CoreSpeakingDataProcessor:
    def __init__(self):
        # 加载环境变量
        load_dotenv()

        # PG数据库配置
        self.pg_config = {
            'host': os.getenv('PG_DB_HOST'),
            'port': int(os.getenv('PG_DB_PORT', 5432)),
            'user': os.getenv('PG_DB_USER'),
            'password': os.getenv('PG_DB_PASSWORD'),
            'database': os.getenv('PG_DB_DATABASE')
        }

        # MySQL数据库配置
        self.mysql_config = {
            'host': os.getenv('MYSQL_HOST'),
            'port': int(os.getenv('MYSQL_PORT', 3306)),
            'user': os.getenv('MYSQL_USERNAME'),
            'password': os.getenv('MYSQL_PASSWORD'),
            'database': os.getenv('MYSQL_DATABASE'),
            'charset': 'utf8mb4'
        }

        # ES配置
        self.es_config = {
            'host': os.getenv('ES_HOST'),
            'port': int(os.getenv('ES_PORT', 9200)),
            'scheme': os.getenv('ES_SCHEME', 'http'),
            'user': os.getenv('ES_USER'),
            'password': os.getenv('ES_PASSWORD')
        }

        self.data = None

    def convert_date_format(self, date_str):
        """将'20250110'格式转换为数据库查询用的格式"""
        try:
            dt = datetime.strptime(date_str, '%Y%m%d')
            return dt.strftime('%Y-%m-%d')
        except ValueError:
            logger.error(f"日期格式错误: {date_str}, 应为'20250110'格式")
            raise

    def get_next_day(self, date_str):
        """获取下一天的日期"""
        try:
            dt = datetime.strptime(date_str, '%Y%m%d')
            next_day = dt + pd.Timedelta(days=1)
            return next_day.strftime('%Y-%m-%d')
        except ValueError:
            logger.error(f"日期格式错误: {date_str}, 应为'20250110'格式")
            raise

    def step1_extract_from_pg(self):
        """步骤一: 从PG数据库筛选核心互动数据"""
        logger.info("步骤一: 开始从PG数据库筛选数据...")

        start_date = self.convert_date_format(START_DATE)
        end_date_next = self.get_next_day(END_DATE)  # 获取结束日期的下一天
        logger.info(f"时间范围: {start_date} 到 {end_date_next} (不含)")

        # 构建查询SQL - 查询8个分表
        all_data = []
        table_names = [f"user_component_play_record_{i}" for i in range(8)]

        for table_name in table_names:
            logger.info(f"正在处理表: {table_name}")

            # 为每个表创建独立的连接，避免事务问题
            try:
                conn = psycopg2.connect(**self.pg_config)
                logger.debug(f"为表 {table_name} 创建数据库连接")
            except Exception as e:
                logger.error(f"为表 {table_name} 创建数据库连接失败: {e}")
                continue

            # 检查当前表是否存在 user_behavior_info 字段
            has_behavior_info = False
            try:
                with conn.cursor() as cur:
                    cur.execute(
                        """
                        SELECT EXISTS (
                            SELECT 1 FROM information_schema.columns 
                            WHERE table_name = %s 
                              AND column_name = 'user_behavior_info'
                        )
                        """,
                        (table_name,)
                    )
                    res = cur.fetchone()
                    has_behavior_info = bool(res[0]) if res else False
                logger.debug(f"表 {table_name} 是否包含 user_behavior_info: {has_behavior_info}")
            except Exception as e:
                logger.warning(f"检测表 {table_name} 的 user_behavior_info 字段失败: {e}")

            # 动态构建查询列
            extra_col = ", user_behavior_info" if has_behavior_info else ""
            sql = f"""
            SELECT
                user_id,
                session_id,
                c_type,
                c_id,
                play_result,
                updated_at{extra_col}
            FROM {table_name}
            WHERE
                updated_at >= %s
                AND updated_at < %s
                AND c_type IN ('core_speaking_reply', 'core_speaking_inquiry')
            ORDER BY updated_at
            """

            try:
                df = pd.read_sql(sql, conn, params=[start_date, end_date_next])
                # 保证列存在，即使部分分表没有该字段
                if 'user_behavior_info' not in df.columns:
                    df['user_behavior_info'] = ''
                if not df.empty:
                    logger.info(f"表 {table_name} 获取到 {len(df)} 条数据")
                    all_data.append(df)
                else:
                    logger.info(f"表 {table_name} 无符合条件的数据")
            except Exception as e:
                logger.error(f"查询表 {table_name} 失败: {e}")
            finally:
                conn.close()

        if all_data:
            self.data = pd.concat(all_data, ignore_index=True)
            logger.info(f"步骤一完成: 总共获取到 {len(self.data)} 条数据")

            # 统计 user_behavior_info 非空条数
            if 'user_behavior_info' in self.data.columns:
                non_empty_behavior = (self.data['user_behavior_info'].astype(str).str.strip() != '').sum()
                logger.info(f"步骤一: user_behavior_info 字段有值 {non_empty_behavior}/{len(self.data)} 条")

            # 处理datetime字段，去掉时区信息（Excel不支持带时区的datetime）
            if 'updated_at' in self.data.columns:
                self.data['updated_at'] = pd.to_datetime(self.data['updated_at']).dt.tz_localize(None)
                logger.info("已处理updated_at字段的时区信息")

            # 输出中间Excel文件
            intermediate_file = f"core_speaking_step1_{START_DATE}_{END_DATE}.xlsx"
            self.data.to_excel(intermediate_file, index=False)
            logger.info(f"步骤一中间文件已保存: {intermediate_file}")
        else:
            logger.warning("步骤一: 未获取到任何数据")
            self.data = pd.DataFrame()

    def step2_add_title_from_mysql(self):
        """步骤二: 从MySQL补充title字段，并从component_config中提取reference_dialog等字段"""
        if self.data.empty:
            logger.warning("步骤二: 数据为空，跳过")
            return

        logger.info("步骤二: 开始从MySQL补充title字段...")

        # 连接MySQL数据库
        try:
            conn = pymysql.connect(**self.mysql_config)
            logger.info("MySQL数据库连接成功")
        except Exception as e:
            logger.error(f"MySQL数据库连接失败: {e}")
            raise

        # 获取所有需要查询的c_type和c_id组合
        unique_components = self.data[['c_type', 'c_id']].drop_duplicates()
        logger.info(f"需要查询 {len(unique_components)} 个不同的组件配置")

        # 查询title、component_config、related_path和kp_relation_info
        sql = """
        SELECT c_type, c_id, title, component_config, related_path, kp_relation_info
        FROM core_interaction_component
        WHERE (c_type, c_id) IN ({})
        """.format(','.join(['(%s,%s)'] * len(unique_components)))

        params = []
        for _, row in unique_components.iterrows():
            params.extend([row['c_type'], row['c_id']])

        try:
            title_df = pd.read_sql(sql, conn, params=params)
            logger.info(f"从MySQL获取到 {len(title_df)} 条组件配置")
        except Exception as e:
            logger.error(f"查询MySQL失败: {e}")
            title_df = pd.DataFrame(columns=['c_type', 'c_id', 'title', 'component_config', 'related_path', 'kp_relation_info'])

        conn.close()

        # 从related_path中解析lesson（lessonIndex）
        def extract_lesson(related_path_str):
            if related_path_str is None or related_path_str == '':
                return ''
            try:
                data = json.loads(related_path_str)
                if isinstance(data, dict):
                    lesson_index = data.get('lessonIndex')
                    return str(lesson_index) if lesson_index is not None else ''
                return ''
            except Exception:
                return ''

        # 从component_config中解析reference_dialog、in_progress_goal和final_goal
        def extract_config_fields(cfg_str):
            result = {
                'reference_dialog': '',
                'in_progress_goal': '',
                'final_goal': ''
            }
            if cfg_str is None or cfg_str == '':
                return result
            try:
                data = json.loads(cfg_str)
                if isinstance(data, dict):
                    dialog_config = data.get('dialogConfig') or data.get('dialog_config')
                    if isinstance(dialog_config, dict):
                        config_obj = dialog_config.get('config')
                        if isinstance(config_obj, dict):
                            promptInfo = config_obj.get('promptInfo')
                            if isinstance(promptInfo, dict):
                                ref = promptInfo.get('reference_dialog')
                                result['reference_dialog'] = ref if isinstance(ref, str) else ''

                                in_prog = promptInfo.get('in_progress_goal')
                                result['in_progress_goal'] = in_prog if isinstance(in_prog, str) else ''

                                final = promptInfo.get('final_goal')
                                result['final_goal'] = final if isinstance(final, str) else ''

                                return result

                    # 兜底：如果顶层就有这些字段
                    ref = data.get('reference_dialog')
                    result['reference_dialog'] = ref if isinstance(ref, str) else ''

                    in_prog = data.get('in_progress_goal')
                    result['in_progress_goal'] = in_prog if isinstance(in_prog, str) else ''

                    final = data.get('final_goal')
                    result['final_goal'] = final if isinstance(final, str) else ''

                return result
            except Exception:
                return result

        # 解析lesson
        if 'related_path' in title_df.columns:
            title_df['lesson'] = title_df['related_path'].apply(extract_lesson)
        else:
            title_df['lesson'] = ''

        # 解析knowledge_points（直接读取kp_relation_info）
        if 'kp_relation_info' in title_df.columns:
            title_df['knowledge_points'] = title_df['kp_relation_info'].fillna('')
        else:
            title_df['knowledge_points'] = ''

        # 解析component_config中的多个字段
        if 'component_config' in title_df.columns:
            config_fields = title_df['component_config'].apply(extract_config_fields)
            title_df['reference_dialog'] = config_fields.apply(lambda x: x['reference_dialog'])
            title_df['in_progress_goal'] = config_fields.apply(lambda x: x['in_progress_goal'])
            title_df['final_goal'] = config_fields.apply(lambda x: x['final_goal'])
        else:
            title_df['reference_dialog'] = ''
            title_df['in_progress_goal'] = ''
            title_df['final_goal'] = ''

        # 仅保留需要合并的列
        title_df = title_df[['c_type', 'c_id', 'title', 'lesson', 'knowledge_points',
                              'in_progress_goal', 'final_goal', 'reference_dialog']]

        # 合并数据
        self.data = pd.merge(
            self.data,
            title_df,
            on=['c_type', 'c_id'],
            how='left'
        )

        # 填充空值
        self.data['title'] = self.data['title'].fillna('')
        self.data['lesson'] = self.data['lesson'].fillna('')
        self.data['knowledge_points'] = self.data['knowledge_points'].fillna('')
        self.data['in_progress_goal'] = self.data['in_progress_goal'].fillna('')
        self.data['final_goal'] = self.data['final_goal'].fillna('')
        self.data['reference_dialog'] = self.data['reference_dialog'].fillna('')

        # 统计解析成功的字段条数
        non_empty_ref = (self.data['reference_dialog'] != '').sum()
        non_empty_lesson = (self.data['lesson'] != '').sum()
        non_empty_kp = (self.data['knowledge_points'] != '').sum()
        non_empty_in_prog = (self.data['in_progress_goal'] != '').sum()
        non_empty_final = (self.data['final_goal'] != '').sum()

        logger.info(f"步骤二完成: 已补充字段统计:")
        logger.info(f"  - lesson: {non_empty_lesson}/{len(self.data)} 条有值")
        logger.info(f"  - knowledge_points: {non_empty_kp}/{len(self.data)} 条有值")
        logger.info(f"  - in_progress_goal: {non_empty_in_prog}/{len(self.data)} 条有值")
        logger.info(f"  - final_goal: {non_empty_final}/{len(self.data)} 条有值")
        logger.info(f"  - reference_dialog: {non_empty_ref}/{len(self.data)} 条有值")

        # 输出中间Excel文件
        intermediate_file = f"core_speaking_step2_{START_DATE}_{END_DATE}.xlsx"
        # 处理datetime字段，去掉时区信息（Excel不支持带时区的datetime）
        if 'updated_at' in self.data.columns:
            self.data['updated_at'] = pd.to_datetime(self.data['updated_at']).dt.tz_localize(None)
        self.data.to_excel(intermediate_file, index=False)
        logger.info(f"步骤二中间文件已保存: {intermediate_file}")

    def step3_add_chat_log_from_es(self):
        """步骤三: 从ES补充对话历史数据"""
        if self.data.empty:
            logger.warning("步骤三: 数据为空，跳过")
            return

        logger.info("步骤三: 开始从ES补充对话历史数据...")

        # 连接ES
        try:
            es_url = f"{self.es_config['scheme']}://{self.es_config['host']}:{self.es_config['port']}"
            if self.es_config['user'] and self.es_config['password']:
                es = Elasticsearch(
                    [es_url],
                    http_auth=(self.es_config['user'], self.es_config['password'])
                )
            else:
                es = Elasticsearch([es_url])

            # 测试连接
            if es.ping():
                logger.info("ES连接成功")
            else:
                raise Exception("ES连接失败")
        except Exception as e:
            logger.error(f"ES连接失败: {e}")
            # 添加空的chat_log和round_num字段
            self.data['chat_log'] = ''
            self.data['round_num'] = 0
            return

        # 获取唯一的session_id
        unique_sessions = self.data['session_id'].unique()
        logger.info(f"需要查询 {len(unique_sessions)} 个不同的session")

        # 批量查询ES
        chat_logs = {}
        round_nums = {}

        batch_size = 100
        for i in range(0, len(unique_sessions), batch_size):
            batch_sessions = unique_sessions[i:i+batch_size]
            logger.info(f"正在处理session批次 {i//batch_size + 1}/{(len(unique_sessions)-1)//batch_size + 1}")

            try:
                # 构建ES查询
                query = {
                    "query": {
                        "bool": {
                            "must": [
                                {"terms": {"session_id": batch_sessions.tolist()}},
                                {"term": {"action": "get_chat"}}
                            ]
                        }
                    },
                    "size": 10000,
                    "sort": [
                        {"session_id": {"order": "asc"}},
                        {"current_round": {"order": "asc"}}
                    ]
                }

                response = es.search(index="llm_roleplayagent_round_log", body=query)
                hits = response['hits']['hits']

                logger.info(f"本批次从ES获取到 {len(hits)} 条对话记录")

                # 按session_id分组处理
                session_rounds = {}
                for hit in hits:
                    source = hit['_source']
                    session_id = source.get('session_id')
                    current_round = source.get('current_round', 0)

                    if session_id not in session_rounds:
                        session_rounds[session_id] = []

                    session_rounds[session_id].append({
                        'current_round': current_round,
                        'user_input': source.get('user_input', ''),
                        'npc_message': source.get('npc_message', '')
                    })

                # 为每个session构建chat_log
                for session_id, rounds in session_rounds.items():
                    # 按round排序
                    rounds.sort(key=lambda x: x['current_round'])

                    chat_parts = []
                    max_round = 0

                    for round_data in rounds:
                        current_round = round_data['current_round']
                        max_round = max(max_round, current_round)

                        if current_round == 0:
                            # round 0 只添加npc_message
                            if round_data['npc_message']:
                                chat_parts.append(f"npc:{round_data['npc_message']}")
                        else:
                            # round 1~n 添加user_input和npc_message
                            if round_data['user_input']:
                                chat_parts.append(f"user:{round_data['user_input']}")
                            if round_data['npc_message']:
                                chat_parts.append(f"npc:{round_data['npc_message']}")

                    chat_logs[session_id] = '\n'.join(chat_parts)
                    round_nums[session_id] = max_round

            except Exception as e:
                logger.error(f"查询ES批次失败: {e}")
                continue

        logger.info(f"完成ES查询，获取到 {len(chat_logs)} 个session的对话记录")

        # 添加chat_log和round_num字段
        self.data['chat_log'] = self.data['session_id'].map(chat_logs).fillna('')
        self.data['round_num'] = self.data['session_id'].map(round_nums).fillna(0)

        logger.info("步骤三完成: 对话历史数据已补充")

    def export_final_excel(self):
        """导出最终Excel文件"""
        if self.data.empty:
            logger.warning("数据为空，无法导出")
            return

        logger.info("开始导出最终Excel文件...")

        # 确保字段顺序
        final_columns = [
            'user_id', 'session_id', 'c_type', 'c_id',
            'play_result', 'updated_at', 'title', 'lesson', 'knowledge_points',
            'in_progress_goal', 'final_goal', 'reference_dialog',
            'chat_log', 'user_behavior_info', 'round_num'
        ]

        # 重新排列列顺序
        self.data = self.data[final_columns]

        # 处理datetime字段，去掉时区信息（Excel不支持带时区的datetime）
        if 'updated_at' in self.data.columns:
            self.data['updated_at'] = pd.to_datetime(self.data['updated_at']).dt.tz_localize(None)
            logger.info("最终导出时已处理updated_at字段的时区信息")

        # 生成文件名
        output_file = f"core_speaking_final_{START_DATE}_{END_DATE}.xlsx"

        # 导出Excel
        self.data.to_excel(output_file, index=False)

        logger.info(f"最终Excel文件已导出: {output_file}")
        logger.info(f"总计导出 {len(self.data)} 条记录")

        # 输出字段统计
        logger.info("字段完整性统计:")
        for col in final_columns:
            if col in ['chat_log', 'title', 'reference_dialog', 'user_behavior_info',
                       'lesson', 'knowledge_points', 'in_progress_goal', 'final_goal']:
                non_empty = (self.data[col] != '').sum()
                logger.info(f"  {col}: {non_empty}/{len(self.data)} 条记录有值")
            elif col == 'round_num':
                non_zero = (self.data[col] > 0).sum()
                logger.info(f"  {col}: {non_zero}/{len(self.data)} 条记录 > 0")

    def process(self):
        """执行完整的数据处理流程"""
        logger.info("="*60)
        logger.info("开始口语核心互动数据处理")
        logger.info(f"时间范围: {START_DATE} - {END_DATE}")
        logger.info("="*60)

        try:
            # 步骤一: PG数据筛选
            self.step1_extract_from_pg()

            # 步骤二: MySQL补充title
            self.step2_add_title_from_mysql()

            # 步骤三: ES补充对话历史
            self.step3_add_chat_log_from_es()

            # 导出最终文件
            self.export_final_excel()

            logger.info("="*60)
            logger.info("数据处理完成!")
            logger.info("="*60)

        except Exception as e:
            logger.error(f"数据处理过程中发生错误: {e}")
            raise

if __name__ == "__main__":
    processor = CoreSpeakingDataProcessor()
    processor.process()
