#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
用户音频数据筛选脚本
功能：从PostgreSQL数据库的分表(user_component_play_record_0~7)中提取指定时间段的用户音频数据。
主要逻辑：
1. 数据源：遍历 user_component_play_record_0 至 user_component_play_record_7 表。
2. 筛选条件：
   - 时间范围：可配置
   - 数据有效性：user_behavior_info 非空且包含 userAudio 和 pronunciationScore。
3. 采样规则：
   - 目标总数：可配置
   - 用户限制：可配置
   - 随机策略：先随机打乱，再按用户分组限制，最后补齐或截断至目标数量。
4. 输出：导出为Excel文件。
   包含字段：
   - index: 序号
   - source_table: 来源表名
   - created_at: 创建时间
   - user_id: 用户ID
   - component_unique_code: 组件唯一标识
   - pronunciationScore: 发音评分
   - userAudio: 音频链接
   - expressContent: 朗读内容文本
"""

import os
import json
import re
import random
import psycopg2
import pymysql
import pandas as pd
from datetime import datetime
from typing import List, Dict, Any
from dotenv import load_dotenv

# 配置参数
CONFIG = {
    # 筛选时间范围
    'START_TIME': '2025-11-10 00:00:00+08:00',
    'END_TIME': '2025-12-10 23:59:59+08:00',

    # 采样参数
    'TARGET_TOTAL': 10000,      # 目标总样本数
    'MAX_PER_USER': 20,        # 单个用户最大样本数
    'TABLE_COUNT': 8,          # 分表数量 (0~N-1)

    # 组件类型过滤
    'C_TYPE_FILTER': 'mid_sentence_dialogue'  # 仅筛选对话互动组件
}

class AudioDataExtractor:
    def __init__(self):
        # 加载环境变量
        load_dotenv()

        # PostgreSQL数据库连接配置
        self.db_config = {
            'host': os.getenv('PG_DB_HOST'),
            'port': os.getenv('PG_DB_PORT'),
            'user': os.getenv('PG_DB_USER'),
            'password': os.getenv('PG_DB_PASSWORD'),
            'database': os.getenv('PG_DB_DATABASE')
        }

        # MySQL数据库连接配置
        self.mysql_config = {
            'host': os.getenv('MYSQL_HOST'),
            'user': os.getenv('MYSQL_USERNAME'),
            'password': os.getenv('MYSQL_PASSWORD'),
            'database': "vala_test",
            'port': int(os.getenv('MYSQL_PORT', 3306)),
            'charset': 'utf8mb4'
        }

        # 分表名称列表
        self.table_names = [f'user_component_play_record_{i}' for i in range(CONFIG['TABLE_COUNT'])]


        # 目标总数
        self.target_total = CONFIG['TARGET_TOTAL']
        # 每个用户最多记录数
        self.max_per_user = CONFIG['MAX_PER_USER']
        
    def get_db_connection(self):
        """获取数据库连接"""
        try:
            conn = psycopg2.connect(**self.db_config)
            return conn
        except Exception as e:
            print(f"数据库连接失败: {e}")
            raise
    
    def extract_audio_info(self, user_behavior_info: str) -> Dict[str, Any]:
        """从user_behavior_info字段中提取音频信息"""
        try:
            behavior_data = json.loads(user_behavior_info)
            if isinstance(behavior_data, list) and len(behavior_data) > 0:
                # 取第一个元素
                data = behavior_data[0]
                if 'userAudio' in data and 'pronunciationScore' in data:
                    return {
                        'userAudio': data.get('userAudio'),
                        'pronunciationScore': data.get('pronunciationScore'),
                        'expressContent': data.get('expressContent')
                    }
        except (json.JSONDecodeError, KeyError, IndexError):
            pass
        return {}
    
    def query_table_data(self, table_name: str) -> List[Dict]:
        """查询单个表的数据"""
        conn = self.get_db_connection()
        cursor = conn.cursor()

        try:
            query = f"""
                SELECT user_id, component_unique_code, c_type, c_id, created_at, user_behavior_info
                FROM {table_name}
                WHERE created_at >= '{CONFIG['START_TIME']}'
                AND created_at <= '{CONFIG['END_TIME']}'
                AND c_type = '{CONFIG['C_TYPE_FILTER']}'
                AND user_behavior_info IS NOT NULL
                AND user_behavior_info != ''
            """

            cursor.execute(query)
            rows = cursor.fetchall()

            results = []
            for row in rows:
                user_id, component_unique_code, c_type, c_id, created_at, user_behavior_info = row

                # 提取音频信息
                audio_info = self.extract_audio_info(user_behavior_info)
                if audio_info and 'userAudio' in audio_info and 'pronunciationScore' in audio_info:
                    results.append({
                        'source_table': table_name,
                        'user_id': user_id,
                        'component_unique_code': component_unique_code,
                        'c_type': c_type,
                        'c_id': c_id,
                        'created_at': created_at,
                        'userAudio': audio_info['userAudio'],
                        'pronunciationScore': audio_info['pronunciationScore'],
                        'expressContent': audio_info.get('expressContent')
                    })

            return results

        finally:
            cursor.close()
            conn.close()

    def get_component_configs(self, data: List[Dict]) -> Dict[str, str]:
        """从MySQL批量获取组件配置信息"""
        # 提取所有unique的(c_type, c_id)组合
        unique_components = set()
        for record in data:
            if 'c_type' in record and 'c_id' in record:
                unique_components.add((record['c_type'], record['c_id']))

        if not unique_components:
            print("没有需要查询的组件")
            return {}

        print(f"正在从MySQL查询 {len(unique_components)} 个组件的配置信息...")

        # 连接MySQL
        try:
            conn = pymysql.connect(**self.mysql_config)
            cursor = conn.cursor()

            # 存储组件配置的字典，key为"c_type-c_id"
            component_configs = {}

            # 批量查询
            for c_type, c_id in unique_components:
                query = """
                SELECT component_config
                FROM middle_interaction_component
                WHERE c_type = %s AND c_id = %s
                """
                cursor.execute(query, (c_type, c_id))
                result = cursor.fetchone()

                if result and result[0]:
                    key = f"{c_type}-{c_id}"
                    component_configs[key] = result[0]

            cursor.close()
            conn.close()

            print(f"成功查询到 {len(component_configs)} 个组件配置")
            return component_configs

        except Exception as e:
            print(f"查询MySQL组件配置失败: {e}")
            return {}

    @staticmethod
    def clean_text(text: str) -> str:
        """清理文本：转小写，去除标点符号和空格"""
        if not text:
            return ""
        # 转小写
        text = text.lower()
        # 去除标点符号和特殊字符，只保留字母和数字
        text = re.sub(r'[^\w\s]', '', text)
        # 去除多余空格
        text = re.sub(r'\s+', '', text)
        return text

    @staticmethod
    def levenshtein_distance(s1: str, s2: str) -> int:
        """计算两个字符串的Levenshtein编辑距离"""
        if len(s1) < len(s2):
            return AudioDataExtractor.levenshtein_distance(s2, s1)

        if len(s2) == 0:
            return len(s1)

        previous_row = range(len(s2) + 1)
        for i, c1 in enumerate(s1):
            current_row = [i + 1]
            for j, c2 in enumerate(s2):
                # 插入、删除、替换的成本
                insertions = previous_row[j + 1] + 1
                deletions = current_row[j] + 1
                substitutions = previous_row[j] + (c1 != c2)
                current_row.append(min(insertions, deletions, substitutions))
            previous_row = current_row

        return previous_row[-1]

    def parse_and_filter_by_config(self, data: List[Dict], component_configs: Dict[str, str]) -> List[Dict]:
        """解析组件配置并筛选question.mode == 'read'的记录"""
        print(f"\n开始根据组件配置筛选数据...")
        print(f"筛选前数据量: {len(data)}")

        filtered_data = []
        skipped_no_config = 0
        skipped_invalid_json = 0
        skipped_wrong_mode = 0

        for record in data:
            c_type = record.get('c_type')
            c_id = record.get('c_id')

            if not c_type or not c_id:
                continue

            # 获取组件配置
            key = f"{c_type}-{c_id}"
            config_str = component_configs.get(key)

            if not config_str:
                skipped_no_config += 1
                continue

            try:
                # 解析JSON配置
                config = json.loads(config_str)

                # 检查question.mode == "read"
                question = config.get('question', {})
                mode = question.get('mode')

                if mode == 'read':
                    # 提取question.content作为refText
                    ref_text = question.get('content', '')
                    record['refText'] = ref_text

                    # 计算编辑距离
                    express_content = record.get('expressContent', '')

                    # 清理文本（去除标点和大小写差异）
                    cleaned_express = self.clean_text(express_content)
                    cleaned_ref = self.clean_text(ref_text)

                    # 计算编辑距离
                    edit_distance = self.levenshtein_distance(cleaned_express, cleaned_ref)
                    record['editDistance'] = edit_distance

                    # 计算相对编辑距离
                    ref_len = len(cleaned_ref)
                    if ref_len > 0:
                        relative_edit_distance = round(edit_distance / ref_len, 4)
                    else:
                        relative_edit_distance = 0
                    record['relativeEditDistance'] = relative_edit_distance

                    filtered_data.append(record)
                else:
                    skipped_wrong_mode += 1

            except (json.JSONDecodeError, AttributeError, TypeError):
                skipped_invalid_json += 1
                continue

        print(f"筛选后数据量: {len(filtered_data)}")
        print(f"  - 缺少配置: {skipped_no_config}")
        print(f"  - 配置解析失败: {skipped_invalid_json}")
        print(f"  - mode不是read: {skipped_wrong_mode}")

        return filtered_data

    def collect_all_data(self) -> List[Dict]:
        """收集所有表的数据"""
        all_data = []

        for table_name in self.table_names:
            print(f"正在查询表: {table_name}")
            try:
                table_data = self.query_table_data(table_name)
                all_data.extend(table_data)
                print(f"表 {table_name} 查询到 {len(table_data)} 条记录")
            except Exception as e:
                print(f"查询表 {table_name} 失败: {e}")
                continue

        print(f"总共收集到 {len(all_data)} 条有效记录")

        if not all_data:
            return []

        # 从MySQL获取组件配置
        component_configs = self.get_component_configs(all_data)

        # 根据组件配置筛选数据（只保留question.mode == "read"的记录）
        filtered_data = self.parse_and_filter_by_config(all_data, component_configs)

        return filtered_data
    
    def random_filter_data(self, data: List[Dict]) -> List[Dict]:
        """随机筛选数据（不按评分分段控制）"""
        # 随机打乱所有数据
        shuffled_data = data.copy()
        random.shuffle(shuffled_data)
        
        print(f"开始随机筛选，总共 {len(shuffled_data)} 条记录")
        return shuffled_data
    
    def apply_user_constraints(self, data: List[Dict]) -> List[Dict]:
        """应用用户约束（每个用户最多2条）"""
        user_records = {}
        
        # 按用户分组
        for record in data:
            user_id = record['user_id']
            if user_id not in user_records:
                user_records[user_id] = []
            user_records[user_id].append(record)
        
        # 每个用户最多选择2条
        final_data = []
        for user_id, records in user_records.items():
            if len(records) <= self.max_per_user:
                final_data.extend(records)
            else:
                # 随机选择2条
                selected = random.sample(records, self.max_per_user)
                final_data.extend(selected)
        
        return final_data
    
    def export_to_excel(self, data: List[Dict], filename: str = 'user_audio_data.xlsx'):
        """导出数据到Excel文件"""
        # 准备导出数据
        export_data = []
        for i, record in enumerate(data):
            # 处理时区问题 - 转换为本地时间字符串
            created_at = record['created_at']
            if hasattr(created_at, 'tz_localize'):
                created_at = created_at.tz_localize(None)
            elif hasattr(created_at, 'replace'):
                created_at = created_at.replace(tzinfo=None)
            
            export_data.append({
                'index': i,
                'source_table': record['source_table'],
                'created_at': created_at,
                'user_id': record['user_id'],
                'component_unique_code': record['component_unique_code'],
                'c_type': record.get('c_type'),
                'c_id': record.get('c_id'),
                'pronunciationScore': record['pronunciationScore'],
                'userAudio': record['userAudio'],
                'expressContent': record.get('expressContent'),
                'refText': record.get('refText'),
                'editDistance': record.get('editDistance'),
                'relativeEditDistance': record.get('relativeEditDistance')
            })
        
        # 创建DataFrame并导出
        df = pd.DataFrame(export_data)
        df.to_excel(filename, index=False)
        print(f"数据已导出到: {filename}")
        print(f"总共导出 {len(export_data)} 条记录")
        
        # 打印统计信息
        self.print_statistics(data)
    
    def print_statistics(self, data: List[Dict]):
        """打印统计信息"""
        print("\n=== 数据统计 ===")
        
        # 评分统计（显示分布情况但不按区间分组）
        scores = [record['pronunciationScore'] for record in data]
        print(f"\n评分统计:")
        print(f"  总记录数: {len(scores)}")
        print(f"  最高分: {max(scores)}")
        print(f"  最低分: {min(scores)}")
        print(f"  平均分: {sum(scores) / len(scores):.2f}")
        
        # 用户分布统计
        user_counts = {}
        for record in data:
            user_id = record['user_id']
            user_counts[user_id] = user_counts.get(user_id, 0) + 1
        
        print(f"\n用户统计:")
        print(f"  总用户数: {len(user_counts)}")
        print(f"  平均每用户记录数: {len(data) / len(user_counts):.2f}")
        
        # 表分布统计
        table_counts = {}
        for record in data:
            table = record['source_table']
            table_counts[table] = table_counts.get(table, 0) + 1
        
        print(f"\n表分布:")
        for table, count in sorted(table_counts.items()):
            print(f"  {table}: {count} 条")
    
    def run(self):
        """运行主流程"""
        print("开始提取用户音频数据...")
        
        # 1. 收集所有数据
        all_data = self.collect_all_data()
        
        if not all_data:
            print("未找到符合条件的数据")
            return
        
        # 2. 随机筛选数据（不按评分分段控制）
        filtered_data = self.random_filter_data(all_data)
        
        # 3. 应用用户约束
        final_data = self.apply_user_constraints(filtered_data)
        
        # 4. 如果数据不足500条，尝试补充
        if len(final_data) < self.target_total:
            print(f"当前数据量 {len(final_data)} 条，少于目标 {self.target_total} 条")
            # 从剩余数据中补充
            used_records = set((r['user_id'], r['component_unique_code'], str(r['created_at'])) for r in final_data)
            available_data = [r for r in all_data if (r['user_id'], r['component_unique_code'], str(r['created_at'])) not in used_records]
            
            needed = self.target_total - len(final_data)
            if len(available_data) >= needed:
                additional = random.sample(available_data, needed)
                final_data.extend(additional)
        
        # 5. 如果超过500条，随机选择500条
        if len(final_data) > self.target_total:
            final_data = random.sample(final_data, self.target_total)
        
        # 6. 导出到Excel
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        filename = f"user_audio_data_{timestamp}.xlsx"
        self.export_to_excel(final_data, filename)

def main():
    extractor = AudioDataExtractor()
    extractor.run()

if __name__ == "__main__":
    main()
