#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""Convert benchmark samples into generic thinker+talker input JSONL."""

from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Any


def read_json(path: Path) -> Any:
    with path.open("r", encoding="utf-8") as f:
        return json.load(f)


def read_jsonl(path: Path) -> list[dict[str, Any]]:
    rows = []
    with path.open("r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if line:
                rows.append(json.loads(line))
    return rows


def write_jsonl(path: Path, rows: list[dict[str, Any]]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as f:
        for row in rows:
            f.write(json.dumps(row, ensure_ascii=False) + "\n")


def abs_or_join(root: Path, maybe_rel: str) -> str:
    p = Path(maybe_rel)
    if p.is_absolute():
        return str(p)
    return str((root / p).resolve())


def parse_echomind_demo_name(path: Path) -> dict[str, str]:
    stem = path.stem
    if stem.endswith("_response"):
        stem = stem[: -len("_response")]
    parts = stem.split("_")
    sample_id = "_".join(parts[:2]) if len(parts) >= 2 else stem
    variant = parts[-1] if parts and parts[-1] in {"target", "neutral", "alternative"} else ""
    middle = parts[2:-1] if variant else parts[2:]
    cue = middle[0] if middle else ""
    topic = "_".join(middle[1:]) if len(middle) > 1 else ""
    return {
        "sample_id": sample_id,
        "cue": cue.replace("_", " "),
        "topic": topic.replace("_", " "),
        "variant": variant,
    }


def echomind_demo_response(args: argparse.Namespace) -> None:
    root = Path(args.root_dir).resolve()
    audio_dir = root / "Audio"
    wavs = sorted(
        p
        for p in audio_dir.glob("*.wav")
        if not p.stem.endswith("_response")
        and (not args.voice_type or any(p.stem.endswith(f"_{v}") for v in args.voice_type))
    )
    out = []
    for idx, wav in enumerate(wavs):
        meta = parse_echomind_demo_name(wav)
        system = (
            "This is an EchoMind demo response-generation sample. "
            "Listen to the user's speech and generate a concise empathetic reply. "
            "Use both spoken content and paralinguistic cues. "
            "The filename metadata is only a weak hint, not a substitute for listening."
        )
        if meta["topic"] or meta["cue"] or meta["variant"]:
            system += (
                f" Metadata hint: topic={meta['topic'] or 'unknown'}, "
                f"vocal_cue={meta['cue'] or 'unknown'}, style_variant={meta['variant'] or 'unknown'}."
            )
        out.append(
            {
                "id": f"echomind_demo_resp_{idx:04d}",
                "orig_id": meta["sample_id"],
                "benchmark": "EchoMind",
                "task": "demo_response",
                "language": args.language,
                "audio_path": str(wav),
                "cue_hint": meta["cue"],
                "topic_hint": meta["topic"],
                "voice_type": meta["variant"],
                "system_prompt": system,
                "user_prompt": (
                    "Listen to the audio and provide one emotionally appropriate spoken response. "
                    "Answer only with the response text."
                ),
            }
        )
    if args.limit:
        out = out[: args.limit]
    write_jsonl(Path(args.out), out)
    print(f"[write] {args.out} rows={len(out)}")


def echomind_demo_mcq_lite(args: argparse.Namespace) -> None:
    root = Path(args.root_dir).resolve()
    audio_dir = root / "MCQ"
    wavs = sorted(p for p in audio_dir.glob("*.wav") if not p.stem.endswith("_response"))
    out = []
    for idx, wav in enumerate(wavs):
        meta = parse_echomind_demo_name(wav)
        out.append(
            {
                "id": f"echomind_demo_mcq_lite_{idx:04d}",
                "orig_id": meta["sample_id"],
                "benchmark": "EchoMind",
                "task": "demo_mcq_lite",
                "language": args.language,
                "audio_path": str(wav),
                "cue_hint": meta["cue"],
                "topic_hint": meta["topic"],
                "voice_type": meta["variant"],
                "system_prompt": (
                    "This is an EchoMind demo vocal-cue perception sample. "
                    "No official multiple-choice options are available in the demo folder, "
                    "so describe the most salient vocal cue or emotional state in a short phrase."
                ),
                "user_prompt": (
                    "Listen to the audio. Identify the most salient vocal cue or emotional state "
                    "in a concise phrase."
                ),
            }
        )
    if args.limit:
        out = out[: args.limit]
    write_jsonl(Path(args.out), out)
    print(f"[write] {args.out} rows={len(out)}")


def uro_single(args: argparse.Namespace) -> None:
    dataset = Path(args.dataset).resolve()
    rows = read_jsonl(dataset)
    data_root = dataset.parent
    out = []
    for idx, item in enumerate(rows):
        sample_id = str(item.get("id") or item.get("sample_id") or f"{idx:04d}")
        target_text = item.get("target_text", item.get("source_text", ""))
        if isinstance(target_text, list):
            target_joined = " / ".join(str(x) for x in target_text)
        else:
            target_joined = str(target_text)
        out.append(
            {
                "id": f"{idx:04d}",
                "orig_id": sample_id,
                "benchmark": "URO-Bench",
                "task": args.dataset_name or dataset.parent.name,
                "language": args.language,
                "audio_path": abs_or_join(data_root, item["source_wav"]),
                "source_text": item.get("source_text", ""),
                "target_text": target_joined,
                "system_prompt": args.system_prompt or "",
                "user_prompt": args.user_prompt
                or "Listen to the input audio and respond directly to the user's spoken request. Do not explain your reasoning.",
            }
        )
    if args.limit:
        out = out[: args.limit]
    write_jsonl(Path(args.out), out)
    print(f"[write] {args.out} rows={len(out)}")


def echomind_asr(args: argparse.Namespace) -> None:
    root = Path(args.root_dir).resolve()
    input_dir = root / "dataset" / f"data_{args.data_type}"
    data = read_json(input_dir / f"script_info_{args.data_type}.json")
    out = []
    for d in data:
        info = d["target_audio_info"]
        out.append(
            {
                "id": f"{d['case_id']}__asr",
                "case_id": d["case_id"],
                "benchmark": "EchoMind",
                "task": "asr",
                "language": args.language,
                "audio_path": str((input_dir / "audio" / info["target_input_audio_file"]).resolve()),
                "source_text": d.get("script", ""),
                "target_text": d.get("script", ""),
                "system_prompt": "You are an ASR engine for this task.",
                "user_prompt": "Please transcribe the speech in the input audio into text. Text must contain only the transcript.",
            }
        )
    if args.limit:
        out = out[: args.limit]
    write_jsonl(Path(args.out), out)
    print(f"[write] {args.out} rows={len(out)}")


def echomind_mcq(args: argparse.Namespace) -> None:
    root = Path(args.root_dir).resolve()
    input_dir = root / "dataset" / f"data_{args.data_type}"
    rows = read_json(input_dir / "MCQ" / args.mcq_file)
    out = []
    for d in rows:
        out.append(
            {
                "id": f"{d['question_id']}__{Path(args.mcq_file).stem}",
                "question_id": d["question_id"],
                "benchmark": "EchoMind",
                "task": f"mcq_{Path(args.mcq_file).stem}",
                "language": args.language,
                "audio_path": str((input_dir / "audio" / d["audio_dir"] / d["audio_name"]).resolve()),
                "question": d.get("question", ""),
                "options_text": d.get("options_text", ""),
                "answer": d.get("answer", ""),
                "system_prompt": (
                    "Determine the correct answer to a multiple-choice question based on the audio clip. "
                    "Analyze the audio and select the most accurate answer without extra explanation."
                ),
                "user_prompt": (
                    f"Question: {d.get('question', '')}\nOptions:\n{d.get('options_text', '')}\n"
                    "Please provide your answer exactly in the format: 'The answer is: [A/B/C/D]'."
                ),
            }
        )
    if args.limit:
        out = out[: args.limit]
    write_jsonl(Path(args.out), out)
    print(f"[write] {args.out} rows={len(out)}")


def echomind_response(args: argparse.Namespace) -> None:
    root = Path(args.root_dir).resolve()
    input_dir = root / "dataset" / f"data_{args.data_type}"
    data = read_json(input_dir / f"script_info_{args.data_type}.json")
    voice_types = [v.strip() for v in args.voice_type.split(",") if v.strip()]
    out = []
    for d in data:
        for voice_type in voice_types:
            key = f"{voice_type}_audio_info"
            if key not in d:
                continue
            info = d[key]
            out.append(
                {
                    "id": f"{d['case_id']}__response__{voice_type}",
                    "case_id": d["case_id"],
                    "voice_type": voice_type,
                    "benchmark": "EchoMind",
                    "task": f"response_{args.system_prompt}",
                    "language": args.language,
                    "audio_path": str((input_dir / "audio" / info[f"{voice_type}_input_audio_file"]).resolve()),
                    "topic": d.get("topic", ""),
                    "spoken_info": d.get("spoken_info", ""),
                    "system_prompt": args.system_prompt_text
                    or (
                        f"The current topic/scenario is: {d.get('topic', '')}. "
                        "Provide a one-turn empathetic response based on the input audio. "
                        "Use both spoken content and voice cues when they are useful."
                    ),
                    "user_prompt": "Listen to the audio and provide a concise, emotionally appropriate one-turn response.",
                }
            )
    if args.limit:
        out = out[: args.limit]
    write_jsonl(Path(args.out), out)
    print(f"[write] {args.out} rows={len(out)}")


def generic_jsonl(args: argparse.Namespace) -> None:
    rows = read_jsonl(Path(args.input))
    out = []
    for idx, item in enumerate(rows):
        audio = item.get(args.audio_field)
        if not audio:
            continue
        out.append(
            {
                "id": str(item.get(args.id_field) or f"{idx:04d}"),
                "benchmark": args.benchmark,
                "task": item.get("task", args.task),
                "language": item.get("language", args.language),
                "audio_path": abs_or_join(Path(args.audio_root).resolve(), audio),
                "source_text": item.get("source_text", item.get("question", "")),
                "target_text": item.get("target_text", item.get("answer", "")),
                "system_prompt": item.get("system_prompt", args.system_prompt or ""),
                "user_prompt": item.get(args.prompt_field, args.user_prompt),
            }
        )
    if args.limit:
        out = out[: args.limit]
    write_jsonl(Path(args.out), out)
    print(f"[write] {args.out} rows={len(out)}")


def main() -> None:
    ap = argparse.ArgumentParser()
    sub = ap.add_subparsers(dest="cmd", required=True)

    p = sub.add_parser("uro-single")
    p.add_argument("--dataset", required=True)
    p.add_argument("--out", required=True)
    p.add_argument("--language", default="en")
    p.add_argument("--dataset-name", default="")
    p.add_argument("--system-prompt", default="")
    p.add_argument("--user-prompt", default="")
    p.add_argument("--limit", type=int, default=0)
    p.set_defaults(func=uro_single)

    p = sub.add_parser("echomind-demo-response")
    p.add_argument("--root-dir", required=True)
    p.add_argument("--out", required=True)
    p.add_argument("--language", default="en")
    p.add_argument("--voice-type", nargs="*", default=["target", "neutral", "alternative"])
    p.add_argument("--limit", type=int, default=0)
    p.set_defaults(func=echomind_demo_response)

    p = sub.add_parser("echomind-demo-mcq-lite")
    p.add_argument("--root-dir", required=True)
    p.add_argument("--out", required=True)
    p.add_argument("--language", default="en")
    p.add_argument("--limit", type=int, default=0)
    p.set_defaults(func=echomind_demo_mcq_lite)

    p = sub.add_parser("echomind-asr")
    p.add_argument("--root-dir", required=True)
    p.add_argument("--data-type", default="synthesis")
    p.add_argument("--out", required=True)
    p.add_argument("--language", default="en")
    p.add_argument("--limit", type=int, default=0)
    p.set_defaults(func=echomind_asr)

    p = sub.add_parser("echomind-mcq")
    p.add_argument("--root-dir", required=True)
    p.add_argument("--data-type", default="synthesis")
    p.add_argument("--mcq-file", required=True)
    p.add_argument("--out", required=True)
    p.add_argument("--language", default="en")
    p.add_argument("--limit", type=int, default=0)
    p.set_defaults(func=echomind_mcq)

    p = sub.add_parser("echomind-response")
    p.add_argument("--root-dir", required=True)
    p.add_argument("--data-type", default="synthesis")
    p.add_argument("--voice-type", default="target,neutral,alternative")
    p.add_argument("--system-prompt", default="enhance")
    p.add_argument("--system-prompt-text", default="")
    p.add_argument("--out", required=True)
    p.add_argument("--language", default="en")
    p.add_argument("--limit", type=int, default=0)
    p.set_defaults(func=echomind_response)

    p = sub.add_parser("generic-jsonl")
    p.add_argument("--input", required=True)
    p.add_argument("--out", required=True)
    p.add_argument("--benchmark", required=True)
    p.add_argument("--task", default="")
    p.add_argument("--language", default="en")
    p.add_argument("--audio-root", default=".")
    p.add_argument("--audio-field", default="audio_path")
    p.add_argument("--id-field", default="id")
    p.add_argument("--prompt-field", default="question")
    p.add_argument("--system-prompt", default="")
    p.add_argument("--user-prompt", default="Listen to the audio and answer directly.")
    p.add_argument("--limit", type=int, default=0)
    p.set_defaults(func=generic_jsonl)

    args = ap.parse_args()
    args.func(args)


if __name__ == "__main__":
    main()