#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""Export aligned inputs for response TTS, talker inference, and e2e Omni.

The query construction pipeline produces three related artifacts:
  1. response_controls.jsonl: response-side Global/Control, used as oracle TTS.
  2. thinker_targets.jsonl: constructed thinker output, used as talker input.
  3. query audio dirs: synthesized user-query audio, used as full-Omni input.

This script keeps qid as the stable join key and writes compact jsonl files for
the next inference stages.
"""

from __future__ import annotations

import argparse
import glob
import json
import os
from pathlib import Path
from typing import Any, Dict, Iterable, List, Optional, Tuple


BASE_DIR = Path("/workspace/echoloc/Dataset/Novel/query_data/v2_2000")
DEFAULT_RESPONSE_CONTROLS = BASE_DIR / "thinker_targets/response_controls.jsonl"
DEFAULT_THINKER_TARGETS = BASE_DIR / "thinker_targets/thinker_targets.jsonl"
DEFAULT_QUERY_QWEN3_DIR = BASE_DIR / "query_audio/qwen3tts_structured_zh_serena_a800"
DEFAULT_QUERY_INDEXTTS_DIR = BASE_DIR / "query_audio/indextts_structured_zh_serena_a800_floatsave"
DEFAULT_OUT_DIR = BASE_DIR / "eval_inputs"


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


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


def load_by_qid(path: Path) -> Dict[str, Dict[str, Any]]:
    out: Dict[str, Dict[str, Any]] = {}
    for row in iter_jsonl(path):
        qid = row.get("qid")
        if qid:
            out[str(qid)] = row
    return out


def _read_instruct_json(path: str) -> Optional[Dict[str, Any]]:
    try:
        with open(path, "r", encoding="utf-8") as f:
            return json.load(f)
    except Exception:
        return None


def meta_language(meta: Dict[str, Any], qid: str) -> str:
    language = str(meta.get("language") or "").strip().lower()
    if language:
        return language
    qid = str(qid or "")
    if qid.startswith("vstyle_en_") or "_en_" in qid:
        return "en"
    return "zh"


def discover_query_audio(qwen3_dir: Path, indextts_dir: Path) -> Dict[str, Dict[str, Any]]:
    """Map qid -> synthesized query audio paths.

    The TTS scripts use line_idx subdirectories, but each Qwen3TTS subdir keeps a
    *_instruct.json containing the original qid/instruct_id. We use that to join
    against response_controls and thinker_targets.
    """
    out: Dict[str, Dict[str, Any]] = {}
    for subdir in sorted(qwen3_dir.glob("*")):
        if not subdir.is_dir():
            continue
        line_idx = subdir.name
        instruct_paths = list(subdir.glob("*_instruct.json"))
        if not instruct_paths:
            continue
        meta = _read_instruct_json(str(instruct_paths[0])) or {}
        qid = meta.get("instruct_id")
        if not qid:
            continue
        language = meta_language(meta, str(qid))

        qwen3_vd = subdir / f"{line_idx}_vd_{language}.wav"
        qwen3_cv = subdir / f"{line_idx}_cv_serena_{language}.wav"
        if language != "zh":
            qwen3_vd_zh = subdir / f"{line_idx}_vd_zh.wav"
            qwen3_cv_zh = subdir / f"{line_idx}_cv_serena_zh.wav"
        else:
            qwen3_vd_zh = qwen3_vd
            qwen3_cv_zh = qwen3_cv
        indextts_wav = indextts_dir / line_idx / f"{line_idx}_indextts_spk-zh_emo-zh_serena.wav"
        indextts_control = indextts_dir / line_idx / f"{line_idx}_indextts_spk-zh_emo-zh_serena_control.json"

        query_audio_path = ""
        query_audio_control_path = ""
        if indextts_wav.exists():
            query_audio_path = str(indextts_wav)
            query_audio_control_path = str(indextts_control) if indextts_control.exists() else ""
        elif qwen3_cv.exists():
            query_audio_path = str(qwen3_cv)
        elif qwen3_vd.exists():
            query_audio_path = str(qwen3_vd)
        elif language != "zh" and qwen3_cv_zh.exists():
            query_audio_path = str(qwen3_cv_zh)
        elif language != "zh" and qwen3_vd_zh.exists():
            query_audio_path = str(qwen3_vd_zh)

        if not query_audio_path:
            continue

        out[str(qid)] = {
            "qid": str(qid),
            "language": language,
            "query_audio_line_idx": line_idx,
            "query_audio_path": query_audio_path,
            "query_audio_control_path": query_audio_control_path,
            "query_qwen3_vd_path": str(qwen3_vd) if qwen3_vd.exists() else "",
            "query_qwen3_cv_path": str(qwen3_cv) if qwen3_cv.exists() else "",
            "query_qwen3_vd_zh_path": str(qwen3_vd_zh) if qwen3_vd_zh.exists() else "",
            "query_qwen3_cv_zh_path": str(qwen3_cv_zh) if qwen3_cv_zh.exists() else "",
            "query_tts_meta": meta,
        }
    return out


def combined_text(target: Dict[str, Any]) -> Tuple[str, str]:
    combined = target.get("combined") or {}
    style = str(combined.get("instruct") or "").strip()
    text = str(combined.get("txt") or "").strip()
    return style, text


def make_talker_row(target: Dict[str, Any], aligned: bool) -> Dict[str, Any]:
    style, text = combined_text(target)
    qid = str(target["qid"])
    return {
        "id": qid,
        "qid": qid,
        "language": target.get("language", "zh"),
        "ability": "constructed_thinker_to_talker",
        "query_type": target.get("query_type"),
        "thinker_style": style,
        "thinker_text": text,
        "need_emochange": target.get("need_emochange"),
        "source_query": target.get("source_query"),
        "aligned_with_query_audio": aligned,
    }


def make_e2e_row(
    target: Dict[str, Any],
    response: Dict[str, Any],
    audio: Dict[str, Any],
) -> Dict[str, Any]:
    style, text = combined_text(target)
    qid = str(target["qid"])
    visible_query = (
        response.get("visible_query")
        or (response.get("source_query_candidate") or {}).get("visible_query")
        or target.get("source_query")
        or {}
    )
    return {
        "id": qid,
        "qid": qid,
        "query_type": target.get("query_type") or response.get("query_type"),
        "language": target.get("language", "zh"),
        "query_text": visible_query.get("text", ""),
        "query_audio_path": audio["query_audio_path"],
        "query_audio_control_path": audio.get("query_audio_control_path", ""),
        "oracle_thinker_style": style,
        "oracle_thinker_text": text,
        "oracle_response_text": response.get("audio_content", text.replace(" <|EMO_CHANGE|> ", "")),
        "response_control": response.get("final_generated_control") or response.get("response_generated_control"),
        "source_event": response.get("source_event"),
    }


def make_manifest_row(
    target: Dict[str, Any],
    response: Dict[str, Any],
    audio: Dict[str, Any],
) -> Dict[str, Any]:
    e2e = make_e2e_row(target, response, audio)
    return {
        **e2e,
        "talker_input_id": target["qid"],
        "response_tts_input_qid": response["qid"],
    }


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--response_controls", type=Path, default=DEFAULT_RESPONSE_CONTROLS)
    ap.add_argument("--thinker_targets", type=Path, default=DEFAULT_THINKER_TARGETS)
    ap.add_argument("--query_qwen3_dir", type=Path, default=DEFAULT_QUERY_QWEN3_DIR)
    ap.add_argument("--query_indextts_dir", type=Path, default=DEFAULT_QUERY_INDEXTTS_DIR)
    ap.add_argument("--out_dir", type=Path, default=DEFAULT_OUT_DIR)
    args = ap.parse_args()

    responses = load_by_qid(args.response_controls)
    targets = load_by_qid(args.thinker_targets)
    query_audio = discover_query_audio(args.query_qwen3_dir, args.query_indextts_dir)

    qids_target = set(targets)
    qids_response = set(responses)
    qids_audio = set(query_audio)
    aligned_qids = sorted(qids_target & qids_response & qids_audio)

    args.out_dir.mkdir(parents=True, exist_ok=True)

    talker_all = [
        make_talker_row(targets[qid], aligned=(qid in aligned_qids))
        for qid in sorted(qids_target)
    ]
    talker_aligned = [row for row in talker_all if row["aligned_with_query_audio"]]
    e2e_rows = [
        make_e2e_row(targets[qid], responses[qid], query_audio[qid])
        for qid in aligned_qids
    ]
    manifest_rows = [
        make_manifest_row(targets[qid], responses[qid], query_audio[qid])
        for qid in aligned_qids
    ]
    response_aligned = [responses[qid] for qid in aligned_qids]

    counts = {
        "response_controls": len(responses),
        "thinker_targets": len(targets),
        "query_audio": len(query_audio),
        "aligned": len(aligned_qids),
        "missing_response_for_target": len(qids_target - qids_response),
        "missing_query_audio_for_target": len(qids_target - qids_audio),
    }

    written = {
        "talker_input_all": write_jsonl(args.out_dir / "talker_input_all.jsonl", talker_all),
        "talker_input_aligned": write_jsonl(args.out_dir / "talker_input_aligned.jsonl", talker_aligned),
        "e2e_input_aligned": write_jsonl(args.out_dir / "e2e_input_aligned.jsonl", e2e_rows),
        "aligned_manifest": write_jsonl(args.out_dir / "aligned_manifest.jsonl", manifest_rows),
        "response_tts_input_aligned": write_jsonl(args.out_dir / "response_tts_input_aligned.jsonl", response_aligned),
    }

    with (args.out_dir / "export_summary.json").open("w", encoding="utf-8") as f:
        json.dump({"counts": counts, "written": written}, f, ensure_ascii=False, indent=2)

    print(json.dumps({"counts": counts, "written": written, "out_dir": str(args.out_dir)}, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()