#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""Export query candidates into control2instruct-like JSON for query-audio TTS."""

from __future__ import annotations

import argparse
import json
import re
from pathlib import Path
from typing import Any, Dict, Iterable, List


DEFAULT_INPUT = "/workspace/echoloc/Dataset/Novel/query_data/v2_2000/query_candidates.jsonl"
DEFAULT_OUT_DIR = "/workspace/echoloc/Dataset/Novel/query_data/v2_2000/query_tts_controls"


def sanitize_id(qid: str) -> str:
    return re.sub(r"[^A-Za-z0-9_.-]+", "_", qid)[:180]


def iter_items(path: Path) -> Iterable[Dict[str, Any]]:
    with path.open("r", encoding="utf-8") as f:
        for line in f:
            if line.strip():
                yield json.loads(line)


def query_voice_to_instruct(item: Dict[str, Any]) -> str:
    qtype = item.get("query_type")
    voice = item.get("query_voice") or {}
    text = item.get("visible_query", {}).get("text", "")
    if qtype == "instruction":
        return "请用标准中性、清晰自然的普通话嗓音，语速适中、音量平稳地读出用户指令文本，不额外加入情绪表演。"

    parts = []
    gender = voice.get("speaker_gender")
    age = voice.get("speaker_age")
    timbre = voice.get("timbre")
    emotion = voice.get("emotion")
    intensity = voice.get("emotion_intensity")
    delivery = voice.get("delivery")
    vad = voice.get("vad") or {}
    gender_map = {"male": "男性", "female": "女性", "neutral": "中性"}
    age_map = {
        "child": "儿童",
        "teen": "青少年",
        "young": "年轻",
        "middle_aged": "中年",
        "elderly": "老年",
    }
    speaker_desc = f"{age_map.get(str(age), str(age) if age else '')}{gender_map.get(str(gender), str(gender) if gender else '')}"
    if age or gender or timbre:
        parts.append(f"请以{speaker_desc or '自然'}说话人的声音，音色{timbre or '自然清晰'}")
    if emotion:
        parts.append(f"表达{emotion}，情绪强度约为{intensity}")
    if vad:
        parts.append(
            f"VAD倾向为 valence={vad.get('valence')}, arousal={vad.get('arousal')}, dominance={vad.get('dominance')}"
        )
    if delivery:
        parts.append(f"具体发声方式：{delivery}")
    parts.append("只朗读用户原话，不添加解释或回应。")
    return "，".join(parts).replace("，，", "，")


def build_control(item: Dict[str, Any]) -> Dict[str, Any]:
    text = item.get("visible_query", {}).get("text", "").strip()
    instruct = query_voice_to_instruct(item)
    return {
        "segments": [
            {
                "instruct": instruct,
                "txt": text,
            }
        ],
        "combined": {
            "instruct": instruct,
            "txt": text,
        },
        "combined_no_speaker": {
            "instruct": instruct,
            "txt": text,
        },
        "meta": {
            "qid": item.get("qid"),
            "query_type": item.get("query_type"),
            "source_event": item.get("source_event"),
            "query_voice": item.get("query_voice"),
        },
    }


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--input", default=DEFAULT_INPUT)
    parser.add_argument("--out_dir", default=DEFAULT_OUT_DIR)
    parser.add_argument("--limit", type=int, default=0)
    parser.add_argument("--overwrite", action="store_true")
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    input_path = Path(args.input)
    out_dir = Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    manifest = out_dir / "manifest.jsonl"
    if args.overwrite and manifest.exists():
        manifest.unlink()

    count = 0
    with manifest.open("a", encoding="utf-8") as mf:
        for item in iter_items(input_path):
            if args.limit > 0 and count >= args.limit:
                break
            qid = item.get("qid", "")
            if not qid:
                continue
            item_dir = out_dir / sanitize_id(qid)
            item_dir.mkdir(parents=True, exist_ok=True)
            control_path = item_dir / "control2instruct.json"
            if control_path.exists() and not args.overwrite:
                count += 1
                continue
            control = build_control(item)
            control_path.write_text(json.dumps(control, ensure_ascii=False, indent=2), encoding="utf-8")
            mf.write(json.dumps({"qid": qid, "query_type": item.get("query_type"), "control_path": str(control_path)}, ensure_ascii=False) + "\n")
            count += 1
    print(f"[DONE] exported={count} out_dir={out_dir}", flush=True)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())