"""
run_labeling.py - 单期或批量 AI 打标脚本
用法:
  单期: python run_labeling.py --ep 4 --model m3
  批量: python run_labeling.py --all --model m3
"""

import sys
sys.stdout.reconfigure(encoding='utf-8')
sys.stderr.reconfigure(encoding='utf-8')

import os
import re
import json
import argparse
from pathlib import Path
from datetime import datetime
from openai import OpenAI
from dotenv import load_dotenv

load_dotenv()

BASE_DIR = Path(__file__).parent.parent
TRANSCRIPTS_DIR = BASE_DIR / "benchmark-set" / "transcripts"
PROMPTS_DIR = BASE_DIR / "prompts"
EXPERIMENTS_DIR = BASE_DIR / "experiments"
GROUND_TRUTH = BASE_DIR / "benchmark-set" / "ground-truth.json"

MODEL_CONFIG = {
    "m3": {
        "base_url": "https://api.minimaxi.com/v1",
        "model_name": "MiniMax-M3",
        "api_key_env": "MINIMAX_API_KEY",
    },
    "deepseek-v4-pro": {
        "base_url": "https://api.deepseek.com",
        "model_name": "deepseek-v4-pro",
        "api_key_env": "DEEPSEEK_API_KEY",
    },
    "mimo-v2.5-pro": {
        "base_url": "https://api.xiaomimimo.com/v1",
        "model_name": "mimo-v2.5-pro",
        "api_key_env": "MIMO_API_KEY",
    },
}

ALL_EPISODES = [3, 4, 5, 7, 10, 11, 12, 13, 14, 15]


def load_prompt(field):
    if field == "narrative":
        return (PROMPTS_DIR / "prompt2_narrative.md").read_text(encoding="utf-8")
    raise ValueError(f"Unknown field: {field}")


def load_transcript(ep):
    pattern = f"ep{ep:02d}_*.md"
    files = list(TRANSCRIPTS_DIR.glob(pattern))
    if not files:
        raise FileNotFoundError(f"No transcript found for ep{ep:02d} in {TRANSCRIPTS_DIR}")
    return files[0].read_text(encoding="utf-8"), files[0].name


def load_ground_truth(ep):
    data = json.loads(GROUND_TRUTH.read_text(encoding="utf-8"))
    for episode in data["episodes"]:
        if episode["ep"] == ep:
            return episode
    return None


def parse_prompt(template, transcript):
    parts = template.split("## USER")
    system_prompt = parts[0].replace("# Prompt 2:叙事结构判别", "").strip()
    user_prompt = parts[1].strip().replace("{transcript}", transcript)
    return system_prompt, user_prompt


def extract_json_from_response(raw: str) -> dict:
    """从模型响应中提取 JSON，兼容推理模型的<think>...输出。"""
    # 先去掉<think>...标签及其内容
    text = re.sub(r'<think>.*?', '', raw, flags=re.DOTALL)
    text = text.strip()
    # 去掉markdown代码块
    text = re.sub(r'^```(?:json)?\s*', '', text)
    text = re.sub(r'\s*```$', '', text)
    text = text.strip()
    # 从第一个 { 开始，到最后一个 } 结束
    first_brace = text.find('{')
    last_brace = text.rfind('}')
    if first_brace != -1 and last_brace != -1 and last_brace >= first_brace:
        json_str = text[first_brace:last_brace + 1]
        return json.loads(json_str)
    # 兜底：直接尝试解析
    return json.loads(text)


def call_model(model_key, system_prompt, user_prompt):
    config = MODEL_CONFIG[model_key]
    client = OpenAI(
        api_key=os.environ[config["api_key_env"]],
        base_url=config["base_url"],
    )
    response = client.chat.completions.create(
        model=config["model_name"],
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt},
        ],
        temperature=0.0,
    )
    raw = response.choices[0].message.content
    return extract_json_from_response(raw)


def run_labeling(ep, model_key):
    transcript, fname = load_transcript(ep)
    template = load_prompt("narrative")
    system_prompt, user_prompt = parse_prompt(template, transcript)
    result = call_model(model_key, system_prompt, user_prompt)
    gt = load_ground_truth(ep)
    ts = datetime.now().strftime("%Y%m%d_%H%M%S")
    out = EXPERIMENTS_DIR / f"{ts}_{model_key}_ep{ep:02d}.json"
    out.write_text(
        json.dumps({"episode": ep, "filename": fname, "result": result, "ground_truth": gt}, ensure_ascii=False, indent=2),
        encoding="utf-8",
    )
    print(f"完成 ep{ep:02d} -> {out.name}")
    return result


def main():
    parser = argparse.ArgumentParser(description="AI 打标脚本")
    parser.add_argument("--ep", type=int, help="单期编号")
    parser.add_argument("--all", action="store_true", help="跑全部")
    parser.add_argument("--model", default="m3", help="模型键名")
    args = parser.parse_args()
    if args.all:
        for ep in ALL_EPISODES:
            run_labeling(ep, args.model)
    elif args.ep:
        run_labeling(args.ep, args.model)
    else:
        parser.print_help()


if __name__ == "__main__":
    main()