"""
从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
期望支持配置  开始 日期 和 结束日期  （可以只配置一个  只配 开始日期  则筛选 >= 开始日期的记录， 只配结束日期 则筛选 <= 结束日期的记录）

输出以下字段内容:

userId
userMsg
userName
soeData
audioUrl
asrStatus
componentId
componentType
dataVersion

"""

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

# 加载环境变量
load_dotenv()

# 配置参数
INDEX_NAME = os.getenv("ES_INDEX", "user-audio")
OUTPUT_FILE = "user_audio_data.xlsx"
START_DATE = "2025-10-15 00:00:00"  # 开始日期，格式: YYYY-MM-DD HH:MM:SS，设为None则不限制
END_DATE = "2025-10-17 00:00:00"  # 结束日期，格式: YYYY-MM-DD HH:MM:SS，设为None则不限制

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

# 采样配置参数
MAX_SAMPLES_PER_USER_MSG = 50  # 每个不重复的userMsg最多采样的数据条数
MAX_SAMPLES_PER_USER_ID = 20    # 每个userId最多采样的数据条数

# 需要导出的字段
EXPORT_FIELDS = [
    "userId",
    "userMsg", 
    "userName",
    "soeData",
    "audioUrl",
    "asrStatus",
    "componentId",
    "componentType",
    "dataVersion",
    "timeStr"
]

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": {
                "timeInt": range_query
            }
        })

    # 如果配置了 userId 列表，则仅选取对应 userId 的数据
    if FILTER_USER_IDS:
        print(f"[DEBUG] 应用 userId 过滤: {FILTER_USER_IDS}")
        must_conditions.append({
            "terms": {
                "userId": FILTER_USER_IDS
            }
        })
    
    # 移除soeData的exists查询，改为在应用层进行更精确的过滤
    # 注释掉原来的soeData exists查询
    # must_conditions.append({
    #     "exists": {
    #         "field": "soeData"
    #     }
    # })
    
    # 构建最终查询
    if must_conditions:
        query = {
            "bool": {
                "must": must_conditions
            }
        }
    else:
        query = {"match_all": {}}
        
    print(f"[DEBUG] 查询条件: {query}")
    
    return {
        "query": query,
        "_source": EXPORT_FIELDS,
        "sort": [{"timeInt": {"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)} 条记录")
        
        # 根据是否配置了 userId 列表决定是否跳过过滤与采样逻辑
        if FILTER_USER_IDS:
            print("[DEBUG] 已配置 userId 列表，跳过过滤与采样逻辑，返回全部匹配数据")
            return all_data
        else:
            # 应用过滤和采样逻辑
            filtered_sampled_data = filter_and_sample_data(all_data)
            return filtered_sampled_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)
    
    # 生成带时间戳的文件名
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    base_name = filename.rsplit('.', 1)[0]
    extension = filename.rsplit('.', 1)[1] if '.' in filename else 'xlsx'
    timestamped_filename = f"{base_name}_{timestamp}.{extension}"
    
    try:
        df.to_excel(timestamped_filename, index=False, engine='openpyxl')
        print(f"数据已导出到: {timestamped_filename}")
        print(f"共导出 {len(data)} 条记录")
    except Exception as e:
        print(f"导出Excel时出错: {e}")

def contains_chinese(text):
    """检测文本是否包含中文字符"""
    if not text:
        return False
    chinese_pattern = re.compile(r'[\u4e00-\u9fff]')
    return bool(chinese_pattern.search(text))

def filter_and_sample_data(data):
    """过滤和采样数据"""
    print(f"[DEBUG] 开始过滤和采样，原始数据量: {len(data)}")
    
    # 第一步：过滤数据
    filtered_data = []
    soe_data_empty_count = 0
    soe_data_not_json_count = 0
    chinese_msg_count = 0
    
    for i, item in enumerate(data):
        # 检查soeData是否存在且以"{"开头
        soe_data = item.get('soeData', '')
        if not soe_data:
            soe_data_empty_count += 1
            if i < 5:  # 只打印前5个样本的详细信息
                print(f"[DEBUG] 样本 {i+1}: soeData为空或不存在")
            continue
            
        if not str(soe_data).strip().startswith('{'):
            soe_data_not_json_count += 1
            if i < 5:  # 只打印前5个样本的详细信息
                print(f"[DEBUG] 样本 {i+1}: soeData不以'{{' 开头，内容: {str(soe_data)[:100]}...")
            continue
            
        # 检查userMsg是否不包含中文
        user_msg = item.get('userMsg', '')
        if contains_chinese(user_msg):
            chinese_msg_count += 1
            if i < 5:  # 只打印前5个样本的详细信息
                print(f"[DEBUG] 样本 {i+1}: userMsg包含中文，内容: {user_msg[:50]}...")
            continue
            
        filtered_data.append(item)
        if i < 5:  # 只打印前5个样本的详细信息
            print(f"[DEBUG] 样本 {i+1}: 通过过滤，userMsg: {user_msg[:50]}...")
    
    print(f"[DEBUG] 过滤统计:")
    print(f"  - soeData为空: {soe_data_empty_count} 条")
    print(f"  - soeData不以'{{' 开头: {soe_data_not_json_count} 条")
    print(f"  - userMsg包含中文: {chinese_msg_count} 条")
    print(f"  - 通过过滤的数据: {len(filtered_data)} 条")
    
    # 第二步：按userMsg分组采样
    user_msg_groups = defaultdict(list)
    for item in filtered_data:
        user_msg = item.get('userMsg', '')
        user_msg_groups[user_msg].append(item)
    
    print(f"[DEBUG] 不重复的userMsg数量: {len(user_msg_groups)}")
    
    # 对每个userMsg组进行采样
    sampled_by_msg = []
    for user_msg, items in user_msg_groups.items():
        # 每个userMsg最多取MAX_SAMPLES_PER_USER_MSG条
        sampled_items = items[:MAX_SAMPLES_PER_USER_MSG]
        sampled_by_msg.extend(sampled_items)
        if len(items) > MAX_SAMPLES_PER_USER_MSG:
            print(f"[DEBUG] userMsg '{user_msg}' 有 {len(items)} 条数据，采样了 {MAX_SAMPLES_PER_USER_MSG} 条")
    
    print(f"[DEBUG] 按userMsg采样后数据量: {len(sampled_by_msg)}")
    
    # 第三步：按userId分组采样
    user_id_groups = defaultdict(list)
    for item in sampled_by_msg:
        user_id = item.get('userId', '')
        user_id_groups[user_id].append(item)
    
    print(f"[DEBUG] 不重复的userId数量: {len(user_id_groups)}")
    
    # 对每个userId组进行采样
    final_sampled_data = []
    for user_id, items in user_id_groups.items():
        # 每个userId最多取MAX_SAMPLES_PER_USER_ID条
        sampled_items = items[:MAX_SAMPLES_PER_USER_ID]
        final_sampled_data.extend(sampled_items)
        if len(items) > MAX_SAMPLES_PER_USER_ID:
            print(f"[DEBUG] userId '{user_id}' 有 {len(items)} 条数据，采样了 {MAX_SAMPLES_PER_USER_ID} 条")
    
    print(f"[DEBUG] 最终采样数据量: {len(final_sampled_data)}")
    
    return final_sampled_data

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']
            soe_data = source.get('soeData', '')
            soe_data_preview = str(soe_data)[:100] if soe_data else 'N/A'
            soe_data_starts_with_brace = str(soe_data).strip().startswith('{') if soe_data else False
            
            print(f"  样本 {i+1}:")
            print(f"    timeInt: {source.get('timeInt', 'N/A')}")
            print(f"    timeStr: {source.get('timeStr', 'N/A')}")
            print(f"    soeData存在: {'是' if soe_data else '否'}")
            print(f"    soeData以{{开头: {'是' if soe_data_starts_with_brace else '否'}")
            print(f"    soeData预览: {soe_data_preview}...")
            print(f"    userMsg: {source.get('userMsg', 'N/A')[:50]}...")
            print(f"    userId: {source.get('userId', 'N/A')}")
        
        # 3. 查询时间范围内的数据（不加soeData过滤）
        time_range_query = {
            "query": {
                "range": {
                    "timeInt": {
                        "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. 查询有soeData的数据总量
        soe_data_query = {
            "query": {
                "exists": {
                    "field": "soeData"
                }
            },
            "size": 0
        }
        response = es_client.search(index=INDEX_NAME, body=soe_data_query)
        soe_data_count = response['hits']['total']
        if isinstance(soe_data_count, dict):
            soe_data_count = soe_data_count['value']
        print(f"[DEBUG] 有soeData字段的数据总量: {soe_data_count}")
        
        # 5. 查询时间范围的实际数据分布
        print(f"[DEBUG] 检查时间字段的实际值范围...")
        agg_query = {
            "query": {"match_all": {}},
            "size": 0,
            "aggs": {
                "time_stats": {
                    "stats": {
                        "field": "timeInt"
                    }
                }
            }
        }
        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_USER_IDS:
        print(f"userId过滤: {FILTER_USER_IDS}")
        print("在配置了 userId 的情况下，将导出匹配用户的全部数据，跳过其他过滤与采样")
    else:
        print(f"过滤条件: soeData非空 且 userMsg不包含中文")
        print(f"采样配置: 每个userMsg最多{MAX_SAMPLES_PER_USER_MSG}条，每个userId最多{MAX_SAMPLES_PER_USER_ID}条")
    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()
