"""
导出  流式语音音频  脚本

v1.0
---
原始数据存储于ES数据库中
索引: llm_realtime_asr_log

es相关配置通过以下环境变量
ES_HOST=xxx
ES_PORT=9200
ES_SCHEME=https
ES_USER=elastic
ES_PASSWORD=xxx  （注意这里可能有特殊符号）

需要配置的内容放置在脚本最开头
开始时间 (8位数字年月日)
截止时间 (8位数字年月日)

仅筛选 时间范围内的数据记录
可以基于 timestamp_int 字段内容进行时间筛选  格式样例:1,769,496,892

正常情况  每个 voice_id 会对应两条记录
可以 以 voice_id为单位
最终  按照每个 voice_id 聚合出以下数据:

asr_prompt （其中一条记录会有这个内容）
result_str （其中一条记录会有这个内容）
timestamp (两条记录都会有，保留最新的一条对应的时间) 格式样例: 2023-12-12 12:12:12
voice_id
audio_url 按以下规则拼接:  https://static.valavala.com/vala_llm/realtime_asr_audio_backup/online/{8位年月日}/{voice_id}.wav     8位年月日 基于 timestamp计算  格式 20260121这种
source （其中一条记录会有这个内容）

最终导出一个excel。
---

"""

import os
from datetime import datetime
import requests
import pandas as pd
from dotenv import load_dotenv
from collections import defaultdict
import urllib3

# ==================== 配置区域 ====================
START_DATE = "20251201"  # 开始日期 (8位数字年月日)
END_DATE = "20260131"    # 结束日期 (8位数字年月日)
# =================================================

# 加载环境变量
load_dotenv()

# ES配置
ES_HOST = os.getenv("ES_HOST")
ES_PORT = int(os.getenv("ES_PORT", "9200"))
ES_SCHEME = os.getenv("ES_SCHEME", "https")
ES_USER = os.getenv("ES_USER", "elastic")
ES_PASSWORD = os.getenv("ES_PASSWORD")
ES_INDEX = "llm_realtime_asr_log"

# 每批处理的数据量
SCROLL_SIZE = 1000
SCROLL_TIMEOUT = "5m"


def timestamp_int_from_date(date_str):
    """将8位日期字符串转换为timestamp_int（秒级时间戳）"""
    dt = datetime.strptime(date_str, "%Y%m%d")
    return int(dt.timestamp())


def format_timestamp(ts):
    """将时间戳转换为格式化字符串"""
    if isinstance(ts, (int, float)):
        return datetime.fromtimestamp(ts).strftime("%Y-%m-%d %H:%M:%S")
    return ts


def generate_audio_url(voice_id, timestamp):
    """生成audio_url"""
    date_str = datetime.fromtimestamp(timestamp).strftime("%Y%m%d")
    return f"https://static.valavala.com/vala_llm/realtime_asr_audio_backup/online/{date_str}/{voice_id}.wav"


def connect_es():
    """测试ES连接"""
    print("正在测试 Elasticsearch 连接...")

    # 禁用SSL警告
    if ES_SCHEME == "https":
        try:
            urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
        except Exception:
            pass

    base_url = f"{ES_SCHEME}://{ES_HOST}:{ES_PORT}"
    auth = (ES_USER, ES_PASSWORD) if ES_USER and ES_PASSWORD else None

    try:
        # 测试连接
        resp = requests.get(
            base_url,
            auth=auth,
            timeout=10,
            verify=False if ES_SCHEME == "https" else True
        )
        resp.raise_for_status()

        print(f"✓ 成功连接到 Elasticsearch: {ES_HOST}:{ES_PORT}")
        return True
    except Exception as e:
        print(f"✗ 连接失败: {e}")
        return False


def query_data(start_date, end_date):
    """查询ES数据"""
    start_ts = timestamp_int_from_date(start_date)
    end_ts = timestamp_int_from_date(end_date) + 86400  # 结束日期加一天，包含当天数据

    print(f"\n开始查询数据...")
    print(f"时间范围: {start_date} 至 {end_date}")
    print(f"时间戳范围: {start_ts} 至 {end_ts}")

    # 禁用SSL警告
    if ES_SCHEME == "https":
        try:
            urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
        except Exception:
            pass

    base_url = f"{ES_SCHEME}://{ES_HOST}:{ES_PORT}"
    search_url = f"{base_url}/{ES_INDEX}/_search"
    headers = {"Content-Type": "application/json"}
    auth = (ES_USER, ES_PASSWORD) if ES_USER and ES_PASSWORD else None

    query = {
        "query": {
            "range": {
                "timestamp_int": {
                    "gte": start_ts,
                    "lt": end_ts
                }
            }
        },
        "sort": [{"timestamp_int": {"order": "asc"}}],
        "size": SCROLL_SIZE
    }

    try:
        # 初始查询（使用scroll）
        params = {"scroll": SCROLL_TIMEOUT}
        response = requests.post(
            search_url,
            headers=headers,
            json=query,
            auth=auth,
            params=params,
            timeout=30,
            verify=False if ES_SCHEME == "https" else True
        )
        response.raise_for_status()
        data = response.json()

        scroll_id = data.get("_scroll_id")
        total_hits = data["hits"]["total"]["value"]

        print(f"✓ 查询完成，共找到 {total_hits} 条记录")

        return data, scroll_id, total_hits

    except Exception as e:
        raise RuntimeError(f"ES查询失败: {e}")


def aggregate_by_voice_id(response, scroll_id, total_hits):
    """按voice_id聚合数据"""
    voice_data = defaultdict(list)
    processed_count = 0

    print("\n开始处理数据...")

    # 禁用SSL警告
    if ES_SCHEME == "https":
        try:
            urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
        except Exception:
            pass

    base_url = f"{ES_SCHEME}://{ES_HOST}:{ES_PORT}"
    scroll_url = f"{base_url}/_search/scroll"
    headers = {"Content-Type": "application/json"}
    auth = (ES_USER, ES_PASSWORD) if ES_USER and ES_PASSWORD else None

    while True:
        hits = response["hits"]["hits"]

        if not hits:
            break

        for hit in hits:
            source = hit["_source"]
            voice_id = source.get("voice_id")

            if voice_id:
                voice_data[voice_id].append(source)

            processed_count += 1

        # 打印进度
        progress = (processed_count / total_hits) * 100
        print(f"\r处理进度: {processed_count}/{total_hits} ({progress:.1f}%)", end="")

        # 获取下一批数据
        try:
            scroll_response = requests.post(
                scroll_url,
                headers=headers,
                json={
                    "scroll": SCROLL_TIMEOUT,
                    "scroll_id": scroll_id
                },
                auth=auth,
                timeout=30,
                verify=False if ES_SCHEME == "https" else True
            )
            scroll_response.raise_for_status()
            response = scroll_response.json()

            # 更新 scroll_id（可能会变化）
            scroll_id = response.get("_scroll_id", scroll_id)

        except Exception as e:
            print(f"\n✗ 获取下一批数据失败: {e}")
            break

    print(f"\n✓ 数据处理完成，共处理 {processed_count} 条记录")
    print(f"✓ 找到 {len(voice_data)} 个唯一的 voice_id")

    # 清理scroll
    try:
        clear_scroll_url = f"{base_url}/_search/scroll"
        requests.delete(
            clear_scroll_url,
            headers=headers,
            json={"scroll_id": [scroll_id]},
            auth=auth,
            timeout=10,
            verify=False if ES_SCHEME == "https" else True
        )
    except Exception:
        pass  # 清理失败不影响结果

    return voice_data


def merge_voice_records(voice_data):
    """合并voice_id的记录，只保留恰好2条记录的"""
    print("\n开始聚合 voice_id 数据...")

    merged_data = []
    valid_count = 0
    invalid_count = 0

    for voice_id, records in voice_data.items():
        # 只处理恰好有2条记录的voice_id
        if len(records) != 2:
            invalid_count += 1
            continue

        valid_count += 1

        # 初始化合并后的数据
        merged_record = {
            "voice_id": voice_id,
            "asr_prompt": None,
            "result_str": None,
            "timestamp": None,
            "source": None,
            "audio_url": None
        }

        # 找出最新的timestamp
        max_timestamp = max(
            records[0].get("timestamp_int", 0),
            records[1].get("timestamp_int", 0)
        )

        # 合并数据
        for record in records:
            if record.get("asr_prompt"):
                merged_record["asr_prompt"] = record["asr_prompt"]
            if record.get("result_str"):
                merged_record["result_str"] = record["result_str"]
            if record.get("source"):
                merged_record["source"] = record["source"]

        # 设置timestamp和audio_url
        merged_record["timestamp"] = format_timestamp(max_timestamp)
        merged_record["audio_url"] = generate_audio_url(voice_id, max_timestamp)

        merged_data.append(merged_record)

    print(f"✓ 聚合完成")
    print(f"  - 有效记录（2条/voice_id）: {valid_count}")
    print(f"  - 无效记录（非2条/voice_id）: {invalid_count}")

    return merged_data


def export_to_excel(data, start_date, end_date):
    """导出到Excel"""
    if not data:
        print("\n警告: 没有数据可导出")
        return

    print(f"\n开始导出数据到 Excel...")

    # 创建DataFrame
    df = pd.DataFrame(data)

    # 调整列顺序
    columns = ["voice_id", "asr_prompt", "result_str", "timestamp", "audio_url", "source"]
    df = df[columns]

    # 生成文件名
    output_dir = "output"
    os.makedirs(output_dir, exist_ok=True)
    filename = f"realtime_asr_export_{start_date}_{end_date}.xlsx"
    filepath = os.path.join(output_dir, filename)

    # 导出Excel
    df.to_excel(filepath, index=False, engine="openpyxl")

    print(f"✓ 数据已导出到: {filepath}")
    print(f"✓ 共导出 {len(df)} 条记录")


def main():
    """主函数"""
    print("=" * 60)
    print("流式语音 ASR 数据导出工具 v1.0")
    print("=" * 60)

    start_time = datetime.now()

    try:
        # 测试ES连接
        if not connect_es():
            raise Exception("无法连接到 Elasticsearch，请检查配置")

        # 查询数据
        response, scroll_id, total_hits = query_data(START_DATE, END_DATE)

        if total_hits == 0:
            print("\n没有找到符合条件的数据")
            return

        # 聚合数据
        voice_data = aggregate_by_voice_id(response, scroll_id, total_hits)

        # 合并记录
        merged_data = merge_voice_records(voice_data)

        # 导出Excel
        export_to_excel(merged_data, START_DATE, END_DATE)

        # 统计耗时
        end_time = datetime.now()
        duration = (end_time - start_time).total_seconds()

        print(f"\n{'=' * 60}")
        print(f"✓ 任务完成! 总耗时: {duration:.2f} 秒")
        print(f"{'=' * 60}")

    except Exception as e:
        print(f"\n✗ 错误: {str(e)}")
        import traceback
        traceback.print_exc()


if __name__ == "__main__":
    main()
