"""
从es中 筛选用户数据

es相关配置通过以下环节变量

ES_HOST=xxx
ES_PORT=9200
ES_SCHEME=https
ES_USER=elastic
ES_PASSWORD=xxx


index: user-audio

脚本思路:

给定 一些过滤参数； 给定导出的excel文件名 （在脚本中以变量方式配置就行）

导出我要的字段内容到一个 excel

过滤字段:
timeStr:  字段内容为str 格式为: 2024-12-31 15:53:19
期望支持配置  开始 日期 和 结束日期  （可以只配置一个  只配 开始日期  则筛选 >= 开始日期的记录， 只配结束日期 则筛选 <= 结束日期的记录）

输出字段内容支持配置:


"""

import os
from datetime import datetime
from dotenv import load_dotenv
from elasticsearch import Elasticsearch
import pandas as pd
import urllib.parse
from collections import defaultdict

# 加载环境变量
load_dotenv()

# 配置参数
INDEX_NAME = "llm_ai_tools_log"
OUTPUT_FILE = "单元挑战用户数据_250906_251024.xlsx"
START_DATE = "2025-09-06 00:00:00"  # 开始日期，格式: YYYY-MM-DD HH:MM:SS，设为None则不限制
END_DATE = "2025-10-24 00:00:00"  # 结束日期，格式: YYYY-MM-DD HH:MM:SS，设为None则不限制

# type字段过滤配置：筛选指定类型的记录，为空则不限制
FILTER_TYPES = ["sent_check_challenge", "speaking_topic_challenge"]

# 可选的 userId 过滤配置：配置为[int, ...] 列表；为空则不限制
FILTER_USER_IDS = []  # 例如: [123, 456]

# 需要导出的字段
EXPORT_FIELDS = [
    "type",
    "question", 
    "user_answer",
    "time_total_ms",
    "score",
    "is_passed",
    "model",
    "write_time_str",
    "write_time_int",
]



def create_es_client():
    """创建Elasticsearch客户端"""
    # 获取环境变量并打印调试信息
    es_host = os.getenv('ES_HOST')
    es_port = os.getenv('ES_PORT', 9200)
    es_scheme = os.getenv('ES_SCHEME', 'https')
    es_user = os.getenv('ES_USER')
    es_password = os.getenv('ES_PASSWORD')
    
    print(f"[DEBUG] ES配置信息:")
    print(f"  ES_HOST: {es_host}")
    print(f"  ES_PORT: {es_port}")
    print(f"  ES_SCHEME: {es_scheme}")
    print(f"  ES_USER: {es_user}")
    print(f"  ES_PASSWORD: {'***已设置***' if es_password else '未设置'}")
    
    # 检查必要的环境变量
    if not es_host:
        raise ValueError("ES_HOST环境变量未设置")
    if not es_user:
        raise ValueError("ES_USER环境变量未设置")
    if not es_password:
        raise ValueError("ES_PASSWORD环境变量未设置")
    
    # URL编码用户名和密码，处理特殊字符
    encoded_user = urllib.parse.quote(es_user, safe='')
    encoded_password = urllib.parse.quote(es_password, safe='')
    
    print(f"[DEBUG] 原始密码包含特殊字符，已进行URL编码")
    
    # 方式1: 使用URL中嵌入认证信息
    host_url_with_auth = f"{es_scheme}://{encoded_user}:{encoded_password}@{es_host}:{es_port}"
    print(f"[DEBUG] 连接URL (带认证): {es_scheme}://{encoded_user}:***@{es_host}:{es_port}")
    
    try:
        # 尝试方式1: URL中嵌入认证
        es_config_1 = {
            'hosts': [host_url_with_auth],
            'verify_certs': False,
            'ssl_show_warn': False,
            'request_timeout': 30,
            'retry_on_timeout': True
        }
        
        print("[DEBUG] 尝试方式1: URL中嵌入认证信息")
        es_client = Elasticsearch(**es_config_1)
        
        # 测试连接
        info = es_client.info()
        print(f"[SUCCESS] 方式1连接成功")
        return es_client
        
    except Exception as e1:
        print(f"[DEBUG] 方式1失败: {e1}")
        
        try:
            # 尝试方式2: 使用basic_auth参数
            host_url = f"{es_scheme}://{es_host}:{es_port}"
            es_config_2 = {
                'hosts': [host_url],
                'basic_auth': (es_user, es_password),
                'verify_certs': False,
                'ssl_show_warn': False,
                'request_timeout': 30,
                'retry_on_timeout': True
            }
            
            print("[DEBUG] 尝试方式2: 使用basic_auth参数")
            es_client = Elasticsearch(**es_config_2)
            
            # 测试连接
            info = es_client.info()
            print(f"[SUCCESS] 方式2连接成功")
            return es_client
            
        except Exception as e2:
            print(f"[DEBUG] 方式2失败: {e2}")
            
            try:
                # 尝试方式3: 使用http_auth参数 (旧版本兼容)
                es_config_3 = {
                    'hosts': [host_url],
                    'http_auth': (es_user, es_password),
                    'verify_certs': False,
                    'ssl_show_warn': False,
                    'request_timeout': 30,
                    'retry_on_timeout': True
                }
                
                print("[DEBUG] 尝试方式3: 使用http_auth参数")
                es_client = Elasticsearch(**es_config_3)
                
                # 测试连接
                info = es_client.info()
                print(f"[SUCCESS] 方式3连接成功")
                return es_client
                
            except Exception as e3:
                print(f"[DEBUG] 方式3失败: {e3}")
                print(f"[ERROR] 所有认证方式都失败了")
                raise e3

def build_query(start_date=None, end_date=None):
    """构建ES查询条件"""
    # 构建基础查询条件
    must_conditions = []
    
    # 添加时间范围条件
    if start_date or end_date:
        range_query = {}
        
        if start_date:
            start_timestamp = int(datetime.strptime(start_date, "%Y-%m-%d %H:%M:%S").timestamp())
            range_query["gte"] = start_timestamp
            print(f"[DEBUG] 开始时间戳: {start_timestamp} (对应 {start_date})")
            
        if end_date:
            end_timestamp = int(datetime.strptime(end_date, "%Y-%m-%d %H:%M:%S").timestamp())
            range_query["lte"] = end_timestamp
            print(f"[DEBUG] 结束时间戳: {end_timestamp} (对应 {end_date})")
            
        must_conditions.append({
            "range": {
                "write_time_int": range_query
            }
        })

    # 如果配置了 userId 列表，则仅选取对应 userId 的数据
    if FILTER_USER_IDS:
        print(f"[DEBUG] 应用 userId 过滤: {FILTER_USER_IDS}")
        must_conditions.append({
            "terms": {
                "userId": FILTER_USER_IDS
            }
        })
    
    # 如果配置了 type 列表，则仅选取对应 type 的数据
    if FILTER_TYPES:
        print(f"[DEBUG] 应用 type 过滤: {FILTER_TYPES}")
        must_conditions.append({
            "terms": {
                "type": FILTER_TYPES
            }
        })
    
    # 构建最终查询
    if must_conditions:
        query = {
            "bool": {
                "must": must_conditions
            }
        }
    else:
        query = {"match_all": {}}
        
    print(f"[DEBUG] 查询条件: {query}")
    
    return {
        "query": query,
        "_source": EXPORT_FIELDS,
        "sort": [{"write_time_int": {"order": "desc"}}]
    }

def fetch_data_from_es(es_client, start_date=None, end_date=None):
    """从ES获取数据"""
    query = build_query(start_date, end_date)
    
    try:
        print(f"[DEBUG] 执行ES查询，使用scroll获取全量数据...")
        
        # 使用scroll API获取全量数据
        scroll_size = 1000  # 每次scroll获取的数据量
        scroll_timeout = '2m'  # scroll超时时间
        
        # 初始化scroll
        query['size'] = scroll_size
        response = es_client.search(
            index=INDEX_NAME, 
            body=query,
            scroll=scroll_timeout
        )
        
        scroll_id = response['_scroll_id']
        hits = response['hits']['hits']
        total_hits = response['hits']['total']
        
        # 获取总数（兼容不同ES版本）
        if isinstance(total_hits, dict):
            total_count = total_hits['value']
        else:
            total_count = total_hits
            
        print(f"[DEBUG] ES中匹配的总记录数: {total_count}")
        
        all_data = []
        batch_count = 1
        
        # 处理第一批数据
        for hit in hits:
            source = hit['_source']
            row = {}
            for field in EXPORT_FIELDS:
                row[field] = source.get(field, "")
            all_data.append(row)
        
        print(f"[DEBUG] 已获取第 {batch_count} 批数据，当前总数: {len(all_data)}")
        
        # 继续scroll获取剩余数据
        while len(hits) == scroll_size:
            batch_count += 1
            response = es_client.scroll(scroll_id=scroll_id, scroll=scroll_timeout)
            scroll_id = response['_scroll_id']
            hits = response['hits']['hits']
            
            for hit in hits:
                source = hit['_source']
                row = {}
                for field in EXPORT_FIELDS:
                    row[field] = source.get(field, "")
                all_data.append(row)
            
            print(f"[DEBUG] 已获取第 {batch_count} 批数据，当前总数: {len(all_data)}")
        
        # 清理scroll
        try:
            es_client.clear_scroll(scroll_id=scroll_id)
        except:
            pass  # 忽略清理错误
        
        print(f"[DEBUG] 从ES获取到数据 {len(all_data)} 条记录")
        return all_data
        
    except Exception as e:
        print(f"查询ES时出错: {e}")
        return []

def export_to_excel(data, filename):
    """导出数据到Excel"""
    if not data:
        print("没有数据可导出")
        return
        
    df = pd.DataFrame(data)
    
    try:
        df.to_excel(filename, index=False, engine='openpyxl')
        print(f"数据已导出到: {filename}")
        print(f"共导出 {len(data)} 条记录")
    except Exception as e:
        print(f"导出Excel时出错: {e}")

def debug_es_data(es_client):
    """调试ES数据，了解实际数据情况"""
    print("\n" + "="*60)
    print("开始调试ES数据...")
    
    try:
        # 1. 查询总数据量
        total_query = {
            "query": {"match_all": {}},
            "size": 0
        }
        response = es_client.search(index=INDEX_NAME, body=total_query)
        total_count = response['hits']['total']
        if isinstance(total_count, dict):
            total_count = total_count['value']
        print(f"[DEBUG] ES索引 '{INDEX_NAME}' 中总数据量: {total_count}")
        
        if total_count == 0:
            print("[ERROR] ES索引中没有任何数据！")
            return
        
        # 2. 查询最近的几条数据，了解数据结构
        sample_query = {
            "query": {"match_all": {}},
            "size": 5,
            "sort": [{"_id": {"order": "desc"}}]
        }
        response = es_client.search(index=INDEX_NAME, body=sample_query)
        hits = response['hits']['hits']
        
        print(f"[DEBUG] 获取到 {len(hits)} 条样本数据:")
        for i, hit in enumerate(hits):
            source = hit['_source']
            
            print(f"  样本 {i+1}:")
            print(f"    write_time_int: {source.get('write_time_int', 'N/A')}")
            print(f"    timeStr: {source.get('timeStr', 'N/A')}")
            print(f"    type: {source.get('type', 'N/A')}")
            print(f"    userId: {source.get('userId', 'N/A')}")
        
        # 3. 查询时间范围内的数据
        time_range_query = {
            "query": {
                "range": {
                    "write_time_int": {
                        "gte": int(datetime.strptime(START_DATE, "%Y-%m-%d %H:%M:%S").timestamp()),
                        "lte": int(datetime.strptime(END_DATE, "%Y-%m-%d %H:%M:%S").timestamp())
                    }
                }
            },
            "size": 0
        }
        response = es_client.search(index=INDEX_NAME, body=time_range_query)
        time_range_count = response['hits']['total']
        if isinstance(time_range_count, dict):
            time_range_count = time_range_count['value']
        print(f"[DEBUG] 时间范围内数据量 ({START_DATE} 到 {END_DATE}): {time_range_count}")
        
        # 4. 查询时间范围的实际数据分布
        print(f"[DEBUG] 检查时间字段的实际值范围...")
        agg_query = {
            "query": {"match_all": {}},
            "size": 0,
            "aggs": {
                "time_stats": {
                    "stats": {
                        "field": "write_time_int"
                    }
                }
            }
        }
        response = es_client.search(index=INDEX_NAME, body=agg_query)
        if 'aggregations' in response:
            stats = response['aggregations']['time_stats']
            min_time = stats.get('min')
            max_time = stats.get('max')
            if min_time and max_time:
                min_date = datetime.fromtimestamp(min_time).strftime("%Y-%m-%d %H:%M:%S")
                max_date = datetime.fromtimestamp(max_time).strftime("%Y-%m-%d %H:%M:%S")
                print(f"    最早时间: {min_date} (时间戳: {min_time})")
                print(f"    最晚时间: {max_date} (时间戳: {max_time})")
        
    except Exception as e:
        print(f"[ERROR] 调试ES数据时出错: {e}")
    
    print("="*60 + "\n")

def main():
    """主函数"""
    print("开始从ES获取单元挑战数据...")
    print(f"索引: {INDEX_NAME}")
    print(f"开始日期: {START_DATE if START_DATE else '不限制'}")
    print(f"结束日期: {END_DATE if END_DATE else '不限制'}")
    if FILTER_TYPES:
        print(f"类型过滤: {FILTER_TYPES}")
    if FILTER_USER_IDS:
        print(f"用户ID过滤: {FILTER_USER_IDS}")
    print("-" * 50)
    
    # 检查.env文件是否存在
    env_file = ".env"
    if not os.path.exists(env_file):
        print(f"[ERROR] {env_file} 文件不存在，请创建并配置ES连接信息")
        print("参考 .env.example 文件进行配置")
        return
    
    print(f"[DEBUG] 找到环境配置文件: {env_file}")
    
    # 创建ES客户端
    try:
        es_client = create_es_client()
    except ValueError as e:
        print(f"[ERROR] 配置错误: {e}")
        print("请检查 .env 文件中的ES配置")
        return
    except Exception as e:
        print(f"[ERROR] 创建ES客户端失败: {e}")
        return
    
    # 测试连接
    try:
        print("[DEBUG] 正在测试ES连接...")
        # ES客户端创建函数中已经包含了连接测试，这里不需要重复测试
        print(f"[SUCCESS] ES连接已建立")
    except Exception as e:
        print(f"[ERROR] ES连接失败: {e}")
        print("\n可能的解决方案:")
        print("1. 检查ES服务是否正常运行")
        print("2. 验证.env文件中的ES_HOST、ES_USER、ES_PASSWORD是否正确")
        print("3. 确认网络连接是否正常")
        print("4. 检查ES用户权限是否足够")
        print("5. 密码中包含特殊字符，已尝试URL编码处理")
        return
    
    # 获取数据
    data = fetch_data_from_es(es_client, START_DATE, END_DATE)
    
    # 导出到Excel
    if data:
        export_to_excel(data, OUTPUT_FILE)
    else:
        print("未获取到任何数据")

if __name__ == "__main__":
    main()
