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

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 = {
    "mimo-v2.5-pro": {
        "base_url": "https://api.xiaomimimo.com/v1",
        "model_name": "mimo-v2.5-pro",
        "api_key_env": "MIMO_API_KEY",
    },
    "deepseek-v4-pro": {
        "base_url": "https://api.deepseek.com",
        "model_name": "deepseek-v4-pro",
        "api_key_env": "DEEPSEEK_API_KEY",
    },
}

ALL_EPISODES = list(range(1, 21))

FIELD_PROMPT_MAP = {
    "narrative": "prompt2_narrative.md",
    "classification": "prompt1_classification.md",
    "opening_hook": "prompt3_opening_hook.md",
}


def load_prompt(field):
    if field not in FIELD_PROMPT_MAP:
        raise ValueError(f"Unknown field: {field}, valid: {list(FIELD_PROMPT_MAP.keys())}")
    return (PROMPTS_DIR / FIELD_PROMPT_MAP[field]).read_text(encoding="utf-8")


def load_transcript(ep):
    pattern = f"ep{ep:03d}_*.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):
    """按 ## SYSTEM / ## USER 分隔符拆解 prompt。
    自动剥离 ## SYSTEM 标签之前的标题行。
    """
    parts = template.split("## USER")
    # parts[0] 是 system 部分，可能包含标题行 + ## SYSTEM 标签
    system_raw = parts[0]
    # 如果有 ## SYSTEM 标签，取它之后的内容；否则去除标题行
    if "## SYSTEM" in system_raw:
        system_prompt = system_raw.split("## SYSTEM", 1)[1].strip()
    else:
        # 没有 ## SYSTEM 标签时，去掉第一行（标题行）作为降级处理
        lines = system_raw.strip().splitlines()
        system_prompt = "\n".join(lines[1:]).strip() if len(lines) > 1 else system_raw.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, field="narrative"):
    transcript, fname = load_transcript(ep)
    template = load_prompt(field)
    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}_{field}_ep{ep:03d}.json"
    out.write_text(
        json.dumps({"episode": ep, "filename": fname, "field": field, "result": result, "ground_truth": gt}, ensure_ascii=False, indent=2),
        encoding="utf-8",
    )
    print(f"完成 ep{ep:03d} [{field}] -> {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="mimo-v2.5-pro", help="模型键名")
    parser.add_argument("--field", default="narrative", choices=["narrative", "classification", "opening_hook"],
                        help="打标字段: narrative(叙事结构) / classification(4分类) / opening_hook(开篇钩子)")
    args = parser.parse_args()
    if args.all:
        for ep in ALL_EPISODES:
            run_labeling(ep, args.model, args.field)
    elif args.ep:
        run_labeling(args.ep, args.model, args.field)
    else:
        parser.print_help()


if __name__ == "__main__":
    main()