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

需求一:
先写一个最简单脚本  实现下面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

环境变量读取:
MYSQL_HOST=xxx
MYSQL_USERNAME=xxx
MYSQL_PASSWORD=xxx
MYSQL_DATABASE=xxx
MYSQL_PORT=xxx
-----------
需求二:
在 PGsql数据库中  筛选数据
数据库相关配置  从.env中读取:
PG_DB_HOST = xxx
PG_DB_PORT = xxx
PG_DB_USER = xxx
PG_DB_PASSWORD = xxx
PG_DB_DATABASE = xxx

读取以下数据表:user_unit_challenge_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）
score_text
question_list
updated_at
category
play_time_seconds （读取 play_time 把ms数据转换为秒 保留整数部分）

导出为excel文件

配置参数直接在脚本开头给出即可

需求三:
需求二中 作为步骤一
本需求为步骤二  基于 步骤一的 文档
进行数据聚合

根据每个unit_id + category 进行分组

统计每个分组下的以下数值:
总记录数量
Perfect数量 (读取 score_text =="Perfect")
Good数量 (读取 score_text =="Good")
Oops数量 (读取 score_text =="Oops")
Perfect率 (Perfect数量 / 总记录数量)
Good率 (Good数量 / 总记录数量)
Oops率 (Oops数量 / 总记录数量)

导出为excel 命名为  步骤一名字_stats.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

# 加载环境变量
load_dotenv()

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

def get_id_2_unit_index():
    # 读取数据库配置
    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]):
        print("Error: Missing database configuration in .env file.")
        print("Ensure MYSQL_HOST, MYSQL_USERNAME, MYSQL_PASSWORD, MYSQL_DATABASE are set.")
        return

    try:
        # 连接数据库
        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
        )

        print(f"Connected to database: {db_host}")

        try:
            with connection.cursor() as cursor:
                # 定义 SQL 语句
                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
                """
                
                print(f"Executing SQL: {sql}")
                
                # 执行查询
                cursor.execute(sql)
                
                # 获取所有结果
                results = cursor.fetchall()
                
                print(f"Total records found: {len(results)}")
                print("-" * 30)
                
                # 打印结果
                print(results)
                id_2_unit_index = {}
                for index, row in enumerate(results):
                    id_2_unit_index[row['id']] = index
                
                print("映射结果:")
                print(id_2_unit_index)


                
                print("-" * 30)
                print("Done.")
                return id_2_unit_index

        finally:
            connection.close()

    except Exception as e:
        print(f"An error occurred: {e}")


def export_unit_challenge_data(start_date, end_date, output_filename):
    """
    从PostgreSQL数据库导出单元挑战数据
    """
    # 读取PostgreSQL数据库配置
    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]):
        print("Error: Missing PostgreSQL database configuration in .env file.")
        print("Ensure PG_DB_HOST, PG_DB_PORT, PG_DB_USER, PG_DB_PASSWORD, PG_DB_DATABASE are set.")
        return

    # 获取 id 到 unit_index 的映射
    print("正在获取 unit_id 映射表...")
    id_2_unit_index = get_id_2_unit_index()
    if not id_2_unit_index:
        print("Error: Failed to get id_2_unit_index mapping.")
        return

    # 转换时间格式: "20250110" -> "2025-01-10 00:00:00"
    start_datetime = datetime.strptime(start_date, "%Y%m%d").strftime("%Y-%m-%d 00:00:00")
    end_datetime = datetime.strptime(end_date, "%Y%m%d").strftime("%Y-%m-%d 00:00:00")

    print(f"时间范围: {start_datetime} 至 {end_datetime}")

    try:
        # 连接PostgreSQL数据库
        connection = psycopg2.connect(
            host=pg_host,
            port=int(pg_port),
            user=pg_user,
            password=pg_password,
            database=pg_database,
            cursor_factory=RealDictCursor
        )

        print(f"已连接到 PostgreSQL 数据库: {pg_host}")

        try:
            with connection.cursor() as cursor:
                # 定义SQL查询
                sql = """
                SELECT
                    user_id,
                    story_id,
                    score_text,
                    question_list,
                    updated_at,
                    category,
                    play_time
                FROM user_unit_challenge_question_result
                WHERE deleted_at IS NULL
                  AND updated_at >= %s
                  AND updated_at < %s
                ORDER BY updated_at ASC
                """

                print(f"执行查询...")

                # 执行查询
                cursor.execute(sql, (start_datetime, end_datetime))

                # 获取所有结果
                results = cursor.fetchall()

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

                # 处理数据
                export_data = []
                for row in results:
                    # 映射 story_id 到 unit_id
                    story_id = row['story_id']
                    unit_id = id_2_unit_index.get(story_id, None)

                    # 转换 play_time (毫秒) 为秒 (整数)
                    play_time_seconds = row['play_time'] // 1000 if row['play_time'] else 0

                    # 移除 updated_at 的时区信息（Excel 不支持带时区的 datetime）
                    updated_at = row['updated_at']
                    if updated_at and hasattr(updated_at, 'replace'):
                        updated_at = updated_at.replace(tzinfo=None)

                    export_data.append({
                        'user_id': row['user_id'],
                        'unit_id': unit_id,
                        'score_text': row['score_text'],
                        'question_list': row['question_list'],
                        'updated_at': updated_at,
                        'category': row['category'],
                        'play_time_seconds': play_time_seconds
                    })

                # 导出到Excel
                if export_data:
                    df = pd.DataFrame(export_data)
                    df.to_excel(output_filename, index=False, engine='openpyxl')
                    print(f"数据已导出到: {output_filename}")
                    print(f"共导出 {len(export_data)} 条记录")
                else:
                    print("没有数据可导出")

        finally:
            connection.close()
            print("数据库连接已关闭")

    except Exception as e:
        print(f"发生错误: {e}")


def aggregate_stats(input_filename):
    """
    基于步骤一的Excel文件进行数据聚合
    按 unit_id + category 分组，统计各项指标
    """
    try:
        # 读取步骤一导出的Excel文件
        print(f"正在读取文件: {input_filename}")
        df = pd.read_excel(input_filename, engine='openpyxl')

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

        # 按 unit_id + category 分组统计
        grouped = df.groupby(['unit_id', 'category'], dropna=False)

        stats_data = []
        for (unit_id, category), group in grouped:
            total_count = len(group)
            perfect_count = (group['score_text'] == 'Perfect').sum()
            good_count = (group['score_text'] == 'Good').sum()
            oops_count = (group['score_text'] == 'Oops').sum()

            # 计算占比
            perfect_rate = round(perfect_count / total_count if total_count > 0 else 0, 2)
            good_rate = round(good_count / total_count if total_count > 0 else 0, 2)
            oops_rate = round(oops_count / total_count if total_count > 0 else 0, 2)

            stats_data.append({
                'unit_id': unit_id,
                'category': category,
                '总记录数量': total_count,
                'Perfect数量': perfect_count,
                'Good数量': good_count,
                'Oops数量': oops_count,
                'Perfect率': perfect_rate,
                'Good率': good_rate,
                'Oops率': oops_rate
            })

        # 生成输出文件名
        base_name = os.path.splitext(input_filename)[0]
        output_filename = f"{base_name}_stats.xlsx"

        # 导出统计结果
        if stats_data:
            stats_df = pd.DataFrame(stats_data)
            stats_df.to_excel(output_filename, index=False, engine='openpyxl')
            print(f"统计数据已导出到: {output_filename}")
            print(f"共 {len(stats_data)} 个分组")
        else:
            print("没有数据可统计")

    except Exception as e:
        print(f"数据聚合时发生错误: {e}")


if __name__ == "__main__":
    # 步骤一：执行导出
    print("=" * 50)
    print("步骤一：导出原始数据")
    print("=" * 50)
    export_unit_challenge_data(START_DATE, END_DATE, OUTPUT_FILENAME)

    # 步骤二：数据聚合
    print("\n" + "=" * 50)
    print("步骤二：数据聚合统计")
    print("=" * 50)
    aggregate_stats(OUTPUT_FILENAME)

    print("\n" + "=" * 50)
    print("全部完成！")
    print("=" * 50)

