"""
** 不要改动我的需求描述，直接在需求后面写代码即可 **

课程巩固  数据导出 和 分析

-----------
需求一:
在 PGsql数据库中  筛选数据
数据库相关配置  从.env中读取:
PG_DB_HOST = xxx
PG_DB_PORT = xxx
PG_DB_USER = xxx
PG_DB_PASSWORD = xxx
PG_DB_DATABASE = xxx

读取以下数据表: user_unit_review_question_result

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

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

在这些时间范围内，筛选数据 (要求deleted_at字段内容为null)

导出以下字段:

user_id
unit_id  （读取每条记录的story_id， 根据 get_id_2_unit_index 函数返回的映射表 映射到 unit_id）
lesson_id （读取chapter_id， 根据该值 查询 mysql表 vala_game_chapter 的 id == chapter_id， 并返回该记录的 index字段的值）
question_list
题目总数
正确数量
正确率
play_time_seconds （读取 play_time 把ms数据转换为秒 保留整数部分）
updated_at

其中 题目总数  正确数量  正确率  都通过  question_list 计算，
该字段为 list of json:
[
  {
    "question": {
      "type": "vocab_meaning_meaning",
      "id": "20-0",
      "title": "“clean” 的意思是什么？",
      "npcId": -1
    },
    "answers": [
      "2"
    ],
    "optionList": [
      {
        "option": "爬行"
      },
      {
        "option": "清晰的"
      },
      {
        "option": "清洁"
      }
    ],
    "isRight": true
  },
  ...
]

每个元素为一道题目， 题目中有 "isRight": true 代表用户做对了。

导出为excel文件
----
需求二  基于 需求一的输出文件  作为 输入文件  进行数据聚合。

聚合的维度是每道题目

根据 question_list 中的 每个题目  取 question -> id 作为唯一标识

统计每个题目  
总记录数量  
正确数量  
正确率

并查询mysql表  补充题目的以下信息:
步骤一中，每个题目id的格式是  num1-num2 (question -> id)
查询vala_kp_question表
其中num1部分 用于 检索vala_kp_question 中的 id,  每个id下 可能有多道题目 在 vala_kp_question的 question 字段 是一个list, num2为question 字段中的索引

补充以下字段:
kp_id (vala_kp_question字段)
category (vala_kp_question字段)
skill (vala_kp_question字段)
type (vala_kp_question字段)
题目配置 (question字段中 对应 num2 索引的内容)

最终针对每道题目输出以下字段:
出现位置 (list,  把所有出现的位置拼接 unit_id +"_"+ lesson_id  例如:"unit10-lesson1" 这样的格式)
question_id (question -> id)
kp_id (vala_kp_question字段)
category (vala_kp_question字段)
skill (vala_kp_question字段)
type (vala_kp_question字段)
题目配置 (question字段中 对应 num2 索引的内容)
总记录数量  
正确数量  
正确率

导出为excel 命名为 步骤一文件_stat.xlsx

所有需要配置的参数  放在脚本开头位置

"""

import os
import pymysql
import psycopg2
from psycopg2.extras import RealDictCursor
from datetime import datetime
import pandas as pd
from dotenv import load_dotenv
import json
from collections import defaultdict

# 加载环境变量
load_dotenv()

# ============ 配置参数 ============
START_DATE = "20250915"  # 起始时间
END_DATE = "20251122"    # 截止时间
OUTPUT_NAME = "lesson_review_data_{}_{}.xlsx".format(START_DATE, END_DATE)  # 输出文件名
OUTPUT_FILENAME = os.path.join("./output", OUTPUT_NAME)
# =================================

def get_mysql_connection():
    """获取MySQL连接"""
    db_host = os.getenv('MYSQL_HOST')
    db_user = os.getenv('MYSQL_USERNAME')
    db_password = os.getenv('MYSQL_PASSWORD')
    db_name = os.getenv('MYSQL_DATABASE')
    db_port = os.getenv('MYSQL_PORT')

    if not all([db_host, db_user, db_password, db_name]):
        raise Exception("Error: Missing MySQL configuration in .env file.")

    connection = pymysql.connect(
        host=db_host,
        user=db_user,
        password=db_password,
        database=db_name,
        port=int(db_port) if db_port else 3306,
        cursorclass=pymysql.cursors.DictCursor
    )
    return connection

def get_pgsql_connection():
    """获取PGsql连接"""
    pg_host = os.getenv('PG_DB_HOST')
    pg_port = os.getenv('PG_DB_PORT')
    pg_user = os.getenv('PG_DB_USER')
    pg_password = os.getenv('PG_DB_PASSWORD')
    pg_database = os.getenv('PG_DB_DATABASE')

    if not all([pg_host, pg_port, pg_user, pg_password, pg_database]):
        raise Exception("Error: Missing PGsql configuration in .env file.")

    connection = psycopg2.connect(
        host=pg_host,
        port=int(pg_port),
        user=pg_user,
        password=pg_password,
        database=pg_database,
        cursor_factory=RealDictCursor
    )
    return connection

def get_id_2_unit_index():
    """获取story_id到unit_id的映射"""
    print("正在获取 story_id 到 unit_id 的映射...")
    connection = get_mysql_connection()

    try:
        with connection.cursor() as cursor:
            sql = """
            SELECT *
            FROM `vala_game_info`
            WHERE id > 0
              AND `vala_game_info`.`deleted_at` IS NULL
            ORDER BY season_package_id asc, `index` asc
            """
            cursor.execute(sql)
            results = cursor.fetchall()

            id_2_unit_index = {}
            for index, row in enumerate(results):
                id_2_unit_index[row['id']] = index

            print(f"成功获取 {len(id_2_unit_index)} 个单元映射")
            return id_2_unit_index
    finally:
        connection.close()

def get_chapter_id_to_lesson_id():
    """获取chapter_id到lesson_id的映射"""
    print("正在获取 chapter_id 到 lesson_id 的映射...")
    connection = get_mysql_connection()

    try:
        with connection.cursor() as cursor:
            sql = """
            SELECT id, `index`
            FROM `vala_game_chapter`
            WHERE deleted_at IS NULL
            """
            cursor.execute(sql)
            results = cursor.fetchall()

            chapter_id_to_lesson_id = {}
            for row in results:
                chapter_id_to_lesson_id[row['id']] = row['index']

            print(f"成功获取 {len(chapter_id_to_lesson_id)} 个课程映射")
            return chapter_id_to_lesson_id
    finally:
        connection.close()

def analyze_question_list(question_list_json):
    """分析题目列表，返回题目总数、正确数量、正确率"""
    try:
        if isinstance(question_list_json, str):
            question_list = json.loads(question_list_json)
        else:
            question_list = question_list_json

        if not isinstance(question_list, list):
            return 0, 0, 0

        total = len(question_list)
        correct = sum(1 for q in question_list if q.get('isRight') == True)
        accuracy = round(correct / total * 100, 2) if total > 0 else 0

        return total, correct, accuracy
    except Exception as e:
        print(f"解析题目列表出错: {e}")
        return 0, 0, 0

def export_step1():
    """需求一：导出原始数据"""
    print("=" * 50)
    print("开始执行需求一：导出原始数据")
    print("=" * 50)

    # 获取映射关系
    id_2_unit_index = get_id_2_unit_index()
    chapter_id_to_lesson_id = get_chapter_id_to_lesson_id()

    # 连接PGsql
    print("正在连接 PGsql 数据库...")
    pg_conn = get_pgsql_connection()

    try:
        with pg_conn.cursor() as cursor:
            # 构建时间范围
            start_datetime = datetime.strptime(START_DATE, "%Y%m%d")
            end_datetime = datetime.strptime(END_DATE, "%Y%m%d")
            end_datetime = end_datetime.replace(hour=23, minute=59, second=59)

            sql = """
            SELECT user_id, story_id, chapter_id, question_list, play_time, updated_at
            FROM user_unit_review_question_result
            WHERE updated_at >= %s
              AND updated_at <= %s
              AND deleted_at IS NULL
            ORDER BY updated_at
            """

            print(f"查询时间范围: {start_datetime} 至 {end_datetime}")
            cursor.execute(sql, (start_datetime, end_datetime))
            results = cursor.fetchall()

            print(f"查询到 {len(results)} 条记录")

            # 处理数据
            export_data = []
            for row in results:
                user_id = row['user_id']
                story_id = row['story_id']
                chapter_id = row['chapter_id']
                question_list_raw = row['question_list']
                play_time = row['play_time']
                updated_at = row['updated_at']

                # 确保 question_list 是 Python 对象（PGsql 的 jsonb 会自动转换）
                # 如果是字符串，先解析；如果已经是对象，直接使用
                if isinstance(question_list_raw, str):
                    try:
                        question_list = json.loads(question_list_raw)
                    except:
                        question_list = []
                else:
                    question_list = question_list_raw if question_list_raw else []

                # 映射 unit_id
                unit_id = id_2_unit_index.get(story_id, -1)

                # 映射 lesson_id
                lesson_id = chapter_id_to_lesson_id.get(chapter_id, -1)

                # 分析题目列表
                total, correct, accuracy = analyze_question_list(question_list)

                # 转换播放时长（ms -> s）
                play_time_seconds = int(play_time / 1000) if play_time else 0

                # 转换question_list为字符串（统一序列化为JSON字符串）
                question_list_str = json.dumps(question_list, ensure_ascii=False) if question_list else ""

                # 移除时区信息（Excel不支持带时区的datetime）
                updated_at_no_tz = updated_at.replace(tzinfo=None) if updated_at else None

                export_data.append({
                    'user_id': user_id,
                    'unit_id': unit_id,
                    'lesson_id': lesson_id,
                    'question_list': question_list_str,
                    '题目总数': total,
                    '正确数量': correct,
                    '正确率': accuracy,
                    'play_time_seconds': play_time_seconds,
                    'updated_at': updated_at_no_tz
                })

            # 导出到Excel
            df = pd.DataFrame(export_data)

            # 确保输出目录存在
            os.makedirs(os.path.dirname(OUTPUT_FILENAME), exist_ok=True)

            df.to_excel(OUTPUT_FILENAME, index=False, engine='openpyxl')
            print(f"成功导出 {len(export_data)} 条记录到: {OUTPUT_FILENAME}")

            return OUTPUT_FILENAME

    finally:
        pg_conn.close()

def get_all_kp_questions(question_ids):
    """批量获取所有题目信息，避免N+1查询问题"""
    print(f"正在批量查询 {len(question_ids)} 道题目的信息...")

    # 解析所有question_id，获取需要查询的kp_question id列表
    kp_ids = set()
    for qid in question_ids:
        try:
            parts = qid.split('-')
            if len(parts) == 2:
                kp_ids.add(int(parts[0]))
        except:
            continue

    print(f"需要查询 {len(kp_ids)} 条 vala_kp_question 记录")

    # 批量查询MySQL
    connection = get_mysql_connection()
    kp_data_map = {}

    try:
        with connection.cursor() as cursor:
            # 使用IN查询批量获取
            if kp_ids:
                placeholders = ','.join(['%s'] * len(kp_ids))
                sql = f"""
                SELECT id, kp_id, category, skill, type, question
                FROM vala_kp_question
                WHERE id IN ({placeholders}) AND deleted_at IS NULL
                """
                cursor.execute(sql, tuple(kp_ids))
                results = cursor.fetchall()

                print(f"成功查询到 {len(results)} 条记录")

                # 构建映射表
                for row in results:
                    kp_data_map[row['id']] = row
    finally:
        connection.close()

    # 为每个question_id构建结果
    question_info_map = {}
    for question_id in question_ids:
        try:
            parts = question_id.split('-')
            if len(parts) != 2:
                question_info_map[question_id] = (None, None, None, None, None)
                continue

            kp_id = int(parts[0])
            question_index = int(parts[1])

            kp_data = kp_data_map.get(kp_id)
            if not kp_data:
                question_info_map[question_id] = (None, None, None, None, None)
                continue

            # 解析question字段
            question_list = kp_data['question']
            if isinstance(question_list, str):
                question_list = json.loads(question_list)

            # 获取指定索引的题目配置
            question_config = None
            if isinstance(question_list, list) and 0 <= question_index < len(question_list):
                question_config = json.dumps(question_list[question_index], ensure_ascii=False)

            question_info_map[question_id] = (
                kp_data['kp_id'],
                kp_data['category'],
                kp_data['skill'],
                kp_data['type'],
                question_config
            )
        except Exception as e:
            print(f"处理题目信息出错 ({question_id}): {e}")
            question_info_map[question_id] = (None, None, None, None, None)

    return question_info_map

def export_step2(input_filename):
    """需求二：数据聚合统计"""
    print("=" * 50)
    print("开始执行需求二：数据聚合统计")
    print("=" * 50)

    # 读取步骤一的输出文件
    print(f"正在读取文件: {input_filename}")
    df = pd.read_excel(input_filename, engine='openpyxl')

    print(f"读取到 {len(df)} 条记录")

    # 按题目聚合统计
    question_stats = defaultdict(lambda: {
        'locations': set(),
        'total_count': 0,
        'correct_count': 0
    })

    parse_success_count = 0
    parse_fail_count = 0
    empty_question_list_count = 0
    processed_question_count = 0

    for idx, row in df.iterrows():
        unit_id = row['unit_id']
        lesson_id = row['lesson_id']
        question_list_str = row['question_list']

        # 解析question_list
        try:
            if pd.isna(question_list_str) or not question_list_str:
                question_list = []
                empty_question_list_count += 1
            else:
                question_list = json.loads(question_list_str)
                parse_success_count += 1
        except Exception as e:
            question_list = []
            parse_fail_count += 1
            if parse_fail_count <= 3:
                print(f"[警告] 第 {idx+1} 条记录解析失败: {e}")

        # 统计每道题目
        for question_item in question_list:
            if not isinstance(question_item, dict):
                continue

            question = question_item.get('question', {})
            question_id = question.get('id')
            is_right = question_item.get('isRight', False)

            if not question_id:
                continue

            # 添加出现位置
            location = f"unit{unit_id}-lesson{lesson_id}"
            question_stats[question_id]['locations'].add(location)

            # 统计数量
            question_stats[question_id]['total_count'] += 1
            if is_right:
                question_stats[question_id]['correct_count'] += 1

            processed_question_count += 1

    print(f"\n解析统计:")
    print(f"  - 解析成功: {parse_success_count} 条")
    print(f"  - 解析失败: {parse_fail_count} 条")
    print(f"  - question_list 为空: {empty_question_list_count} 条")
    print(f"  - 处理的题目总数: {processed_question_count} 道")
    print(f"  - 聚合得到不同题目: {len(question_stats)} 道")

    # 批量获取所有题目信息（优化性能）
    all_question_ids = list(question_stats.keys())
    question_info_map = get_all_kp_questions(all_question_ids)

    # 构建导出数据
    print(f"\n正在构建导出数据...")
    export_data = []
    for idx, (question_id, stats) in enumerate(question_stats.items()):
        if (idx + 1) % 100 == 0:
            print(f"  已处理 {idx + 1}/{len(question_stats)} 道题目")

        # 从批量查询结果中获取题目信息
        kp_id, category, skill, type_field, question_config = question_info_map.get(
            question_id, (None, None, None, None, None)
        )

        # 计算正确率
        total = stats['total_count']
        correct = stats['correct_count']
        accuracy = round(correct / total * 100, 2) if total > 0 else 0

        # 出现位置列表
        locations_list = sorted(list(stats['locations']))
        locations_str = ', '.join(locations_list)

        export_data.append({
            '出现位置': locations_str,
            'question_id': question_id,
            'kp_id': kp_id,
            'category': category,
            'skill': skill,
            'type': type_field,
            '题目配置': question_config,
            '总记录数量': total,
            '正确数量': correct,
            '正确率': accuracy
        })

    # 导出到Excel
    output_stat_filename = input_filename.replace('.xlsx', '_stat.xlsx')
    df_stat = pd.DataFrame(export_data)

    print(f"\n正在导出到 Excel...")
    df_stat.to_excel(output_stat_filename, index=False, engine='openpyxl')

    print(f"成功导出 {len(export_data)} 道题目的统计数据到: {output_stat_filename}")

    return output_stat_filename

def main():
    """主函数"""
    try:
        # 执行需求一
        step1_output = export_step1()

        print("\n")

        # 执行需求二
        step2_output = export_step2(step1_output)

        print("\n" + "=" * 50)
        print("所有任务完成！")
        print(f"需求一输出文件: {step1_output}")
        print(f"需求二输出文件: {step2_output}")
        print("=" * 50)

    except Exception as e:
        print(f"执行出错: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    main()



