import os
import json
import time
import re
from openai import OpenAI
from tqdm import tqdm
from concurrent.futures import ThreadPoolExecutor, as_completed

# ================= 配置区域 =================
API_KEY = os.environ.get("ECHOLOC_API_KEY", "") # 替换为实际KEY
BASE_URL = os.environ.get("LLM_BASE_URL", "")
MODEL_NAME = "deepseek-chat"
MAX_WORKERS = 16       # 并发线程数
MAX_API_RETRIES = 3   # 接口请求失败时的网络重试次数
MAX_EVAL_ROUNDS = 5   # 如果分数低于4，最多允许重新生成的次数

def contains_chinese(text):
    return bool(re.search(r'[\u4e00-\u9fff]', text))

def evaluate_instruct(client, global_desc, emotion_list, generated_instruct, is_zh):
    """
    调用 LLM 对生成的指令进行 1-5 评分，并输出 JSON 格式的反馈
    """
    emotion_flow_str = " -> ".join(emotion_list)
    speaker_info = f"- 说话人特征：{global_desc}" if global_desc else "- 说话人特征：无（仅评估情绪纯度及演变）"
    
    system_prompt = "你是一个极其严苛的 TTS 数据集质检专家。你的任务是评估生成的 TTS 提示词，并严格输出 JSON 格式。"
    
    user_prompt = f"""
请对下方的【生成指令】进行评估打分 (1-5分)。

# 原始输入数据
{speaker_info}
- 目标情绪设定：{emotion_flow_str}

# 待评估的生成指令
{generated_instruct}

# 评分标准
5分：极其自然连贯，高度精炼，准确传达了情绪且没有任何诸如“的语气”、“依次为”之类的冗余/机械字眼。如果没有要求说话人特征，则指令中绝对没有脑补的音色描述。
4分：自然且准确，但可能稍微有一两个多余的修饰词，整体可以直接用于模型训练。
3分：勉强可用，存在明显的机械拼接痕迹（如“A转B转C”），或者遗漏了关键情绪。
1-2分：完全不通顺，语言错误，严重偏离了原始情绪，或错误地加入了未提供的音色特征。

# 输出要求
必须严格输出纯 JSON 对象，不要包含 ```json 的 Markdown 标记，格式如下：
{{
    "score": 4,
    "feedback": "具体的扣分原因以及下一步生成的修改建议（如果满分则填无）。"
}}
"""
    for attempt in range(MAX_API_RETRIES):
        try:
            response = client.chat.completions.create(
                model=MODEL_NAME,
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_prompt}
                ],
                temperature=0.1, # 极低温度，保证打分客观严谨
                response_format={"type": "json_object"} # 强制要求返回 JSON
            )
            result_str = response.choices[0].message.content.strip()
            
            # 容错解析
            result_str = result_str.strip('`').removeprefix('json').strip()
            eval_result = json.loads(result_str)
            
            score = int(eval_result.get("score", 5))
            feedback = eval_result.get("feedback", "")
            return score, feedback
        except Exception as e:
            if attempt == MAX_API_RETRIES - 1:
                return 4, "" # 如果裁判模型一直报错，默认给过，防止阻塞
            time.sleep(2)

def generate_concise_instruct(client, global_desc, emotion_list, is_zh, previous_feedback=None):
    """
    调用 LLM 生成精炼指令。如果传入了 previous_feedback，则要求大模型反思修改。
    """
    emotion_flow_str = " -> ".join(emotion_list)
    target_lang = "中文" if is_zh else "英文"
    
    system_prompt = "你是一个顶级的文本转语音 (TTS) 提示词工程师。"
    
    # 动态构建 User Prompt
    prompt_lines = [f"请将以下信息融合成【一句话】的精炼指令。"]
    
    if global_desc:
        # 有音色信息（对应 N+1 中的 Combined）
        prompt_lines.append(f"\n# 输入数据\n- 说话人基础音色：{global_desc}\n- 情绪序列：{emotion_flow_str}")
        prompt_lines.append("\n# 要求：将音色与情绪平滑融合。不要用“转变”等词描述过程，写成具有戏剧张力的单一复合描述！极度精简！")
    else:
        prompt_lines.append(f"\n# 输入数据\n- 情绪序列：{emotion_flow_str}")
        if len(emotion_list) > 1:
            # 无音色，但有多段情绪（对应 N+2 中的 Combined 无 Speaker）
            prompt_lines.append("\n# 要求：将这一连串情绪平滑融合为一句具有戏剧张力的复合情感演变描述。注意：【绝对不要】加入性别、年龄等任何说话人音色特征！极度精简！")
        else:
            # 无音色，单段情绪（对应 N 个 Segment）
            prompt_lines.append("\n# 要求：极度提纯！去掉“的语气”等所有冗余字眼，直接保留最核心的情感状态。注意：【绝对不要】加入说话人音色特征！")

    prompt_lines.append(f"语言必须是【{target_lang}】。只输出指令本身，不带引号。")

    # 【核心】：注入上一轮裁判的批评意见
    if previous_feedback:
        prompt_lines.append(f"\n# ⚠️ 裁判反馈 (前次生成未达标) ⚠️\n请严格根据以下建议进行反思和重写：\n{previous_feedback}")

    user_prompt = "\n".join(prompt_lines)

    for attempt in range(MAX_API_RETRIES):
        try:
            response = client.chat.completions.create(
                model=MODEL_NAME,
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_prompt}
                ],
                temperature=0.4 if previous_feedback else 0.3, # 如果是重写，稍微提高一点温度寻找新思路
                max_tokens=150
            )
            return response.choices[0].message.content.strip()
        except Exception as e:
            if attempt == MAX_API_RETRIES - 1:
                return None
            time.sleep(2)

def generate_with_eval(client, global_desc, emotion_list, is_zh):
    """
    封装了 [生成 -> 评估 -> 修正] 的闭环核心逻辑
    """
    current_feedback = None
    best_instruct = ""
    best_score = -1

    for round_idx in range(MAX_EVAL_ROUNDS + 1):
        # 1. 生成指令
        instruct = generate_concise_instruct(client, global_desc, emotion_list, is_zh, current_feedback)
        if not instruct:
            break
            
        # 2. 裁判打分
        score, feedback = evaluate_instruct(client, global_desc, emotion_list, instruct, is_zh)
        
        # 记录历史最佳
        if score > best_score:
            best_score = score
            best_instruct = instruct
            
        # 3. 判断是否通关
        if score >= 4:
            break
        else:
            current_feedback = feedback
            print(f"instruct:{instruct}\nscore:{score}\nfeedback:{feedback}", flush=True)

    return best_instruct

def process_single_file(file_path, client):
    """处理单个文件，调用带评估机制的生成方法"""
    with open(file_path, 'r', encoding='utf-8') as f:
        try:
            data = json.load(f)
        except json.JSONDecodeError:
            return False

    control_info = data.get("Control", [])
    if not control_info:
        return False

    full_text = "".join([segment.get("sample_text", "") for segment in control_info])
    if not full_text.strip():
        return False
    
    is_zh = contains_chinese(full_text)
    global_info = data.get("Global", {})
    speaker_desc = global_info.get("instruct_zh", "") if is_zh else global_info.get("instruct_en", "")
    
    # 【结构更新】：新增 combined_no_speaker 字段
    output_data = {"segments": [], "combined": {}, "combined_no_speaker": {}}
    emotion_list = []

    # 1. 处理每个子段落 (不带speaker) —— 共有 N 个
    for segment in control_info:
        txt = segment.get("sample_text", "")
        emo = segment.get("instruct_zh", "") if is_zh else segment.get("instruct_en", "")
        emotion_list.append(emo)
        
        single_instruct = generate_with_eval(client, "", [emo], is_zh)
        output_data["segments"].append({"instruct": single_instruct, "txt": txt})

    # 2. 处理合并段落 (带speaker) —— 第 N+1 个
    combined_instruct = generate_with_eval(client, speaker_desc, emotion_list, is_zh)
    if not combined_instruct:
        combined_instruct = f"{speaker_desc} {', '.join(emotion_list)}"
    output_data["combined"] = {"instruct": combined_instruct, "txt": full_text}

    # 3. 处理合并段落 (不带speaker) —— 第 N+2 个
    combined_no_speaker_instruct = generate_with_eval(client, "", emotion_list, is_zh)
    if not combined_no_speaker_instruct:
        combined_no_speaker_instruct = " -> ".join(emotion_list)
    output_data["combined_no_speaker"] = {"instruct": combined_no_speaker_instruct, "txt": full_text}

    # 4. 保存文件
    dir_name = os.path.dirname(file_path)
    output_path = os.path.join(dir_name, "control2instruct.json")
    try:
        with open(output_path, 'w', encoding='utf-8') as f:
            json.dump(output_data, f, ensure_ascii=False, indent=4)
        return True
    except IOError:
        return False

def main():
    TARGET_DIRS = [
        "/workspace/echoloc/codes/TTS-Framework/logs/vstyle/qwen3tts/ds3_vstyle_controls_3_17",
        "/workspace/echoloc/codes/TTS-Framework/logs/vstyle/qwen3tts/ds3_vad_textrefine_emochange",
        "/workspace/echoloc/codes/TTS-Framework/logs/vstyle/qwen3tts/ds3_vad_resp_control_v1_evaluated_t0.2",
        "/workspace/echoloc/codes/TTS-Framework/logs/vstyle/qwen3tts/ds3_vad_resp_control_v1_evaluated_t0.2_refined",
        "/workspace/echoloc/codes/TTS-Framework/logs/vstyle/indextts/ds3_vstyle_controls_3_17",
        "/workspace/echoloc/codes/TTS-Framework/logs/vstyle/indextts/ds3_vad_resp_control_v1_evaluated_t0.2",
        "/workspace/echoloc/codes/TTS-Framework/logs/vstyle/indextts/ds3_vad_resp_control_v1_evaluated_t0.2_refined",
        "/workspace/echoloc/codes/TTS-Framework/logs/vstyle/indextts/ds3_vad_textrefine_emochange"
    ]
    
    client = OpenAI(api_key=API_KEY, base_url=BASE_URL)
    
    print(f"正在扫描 {len(TARGET_DIRS)} 个主文件夹...")
    json_files = []
    skipped_count = 0
    
    for target_dir in TARGET_DIRS:
        if not os.path.exists(target_dir):
            print(f"⚠️ 警告: 路径不存在，已跳过 -> {target_dir}")
            continue
            
        for root, dirs, files in os.walk(target_dir):
            if "control.json" in files:
                if "control2instruct.json" in files:
                    skipped_count += 1
                    continue
                json_files.append(os.path.join(root, "control.json"))
                
    total_files = len(json_files)
    
    if skipped_count > 0:
        print(f"⏩ 跳过 {skipped_count} 个已完成文件。")
    if total_files == 0:
        print("✅ 无新文件需处理。")
        return

    print(f"🚀 开始处理 {total_files} 个文件 (已启用智能打分与重写机制)...")
    
    success_count = 0
    with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
        future_to_file = {executor.submit(process_single_file, path, client): path for path in json_files}
        
        for future in tqdm(as_completed(future_to_file), total=total_files, desc="Processing"):
            if future.result():
                success_count += 1
                
    print(f"\n🎉 跑批完成！本次成功生成了 {success_count} 个高质量的 control2instruct.json 文件。")

if __name__ == "__main__":
    main()