#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""Prepare and assemble iter006 data from iter005 Chinese assets plus new English assets."""

from __future__ import annotations

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


DEFAULT_ITER005 = Path("/workspace/echoloc/Dataset/Novel/query_data/supervisor_iterations/iter_005")
DEFAULT_ITER006 = Path("/workspace/echoloc/Dataset/Novel/query_data/supervisor_iterations/iter_006")


def read_jsonl(path: Path) -> List[Dict[str, Any]]:
    rows: List[Dict[str, Any]] = []
    if not path.exists():
        return rows
    with path.open("r", encoding="utf-8") as handle:
        for line in handle:
            line = line.strip()
            if line:
                try:
                    rows.append(json.loads(line))
                except json.JSONDecodeError:
                    rows.append(json.loads(line.lstrip("\ufeff；;")))
    return rows


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 handle:
        for row in rows:
            handle.write(json.dumps(row, ensure_ascii=False) + "\n")
            n += 1
    return n


def qid(row: Dict[str, Any]) -> str:
    return str(row.get("qid") or row.get("id") or "")


def is_vstyle_en(row: Dict[str, Any]) -> bool:
    cur = qid(row)
    if cur.startswith("vstyle_en_"):
        return True
    language = str(row.get("language") or ((row.get("source") or {}).get("language")) or "").lower()
    return language == "en" and cur.startswith("vstyle_")


def is_vstyle_zh(row: Dict[str, Any]) -> bool:
    cur = qid(row)
    if cur.startswith("vstyle_en_"):
        return False
    if cur.startswith("vstyle_zh_"):
        return True
    language = str(row.get("language") or ((row.get("source") or {}).get("language")) or "").lower()
    return language == "zh" and cur.startswith("vstyle_")


def is_story_zh(row: Dict[str, Any]) -> bool:
    cur = qid(row)
    if cur.startswith("vstyle_") or cur.startswith("story_en_"):
        return False
    if "vstyle_en_" in cur:
        return False
    language = str(row.get("language") or ((row.get("source") or {}).get("language")) or "zh").lower()
    return language == "zh"


def is_zh_reuse(row: Dict[str, Any]) -> bool:
    return is_story_zh(row) or is_vstyle_zh(row)


def mark_language(row: Dict[str, Any], language: str) -> Dict[str, Any]:
    row = dict(row)
    row["language"] = language
    source = row.get("source")
    if isinstance(source, dict):
        source = dict(source)
        source["language"] = language
        row["source"] = source
    source_candidate = row.get("source_query_candidate")
    if isinstance(source_candidate, dict):
        source_candidate = dict(source_candidate)
        source_candidate["language"] = language
        row["source_query_candidate"] = source_candidate
    return row


def by_qid(rows: Iterable[Dict[str, Any]]) -> Dict[str, Dict[str, Any]]:
    out: Dict[str, Dict[str, Any]] = {}
    for row in rows:
        cur = qid(row)
        if cur:
            out[cur] = row
    return out


def read_first_jsonl(paths: Iterable[Path]) -> List[Dict[str, Any]]:
    for path in paths:
        rows = read_jsonl(path)
        if rows:
            return rows
    return []


def copy_if_exists(src: Path, dst: Path) -> bool:
    if not src.exists():
        return False
    dst.parent.mkdir(parents=True, exist_ok=True)
    shutil.copy2(src, dst)
    return True


def extract_vstyle_en(iter005: Path, iter006: Path) -> Dict[str, Any]:
    out_dir = iter006 / "english_rebuild" / "vstyle_en"
    out_dir.mkdir(parents=True, exist_ok=True)

    candidates = [mark_language(row, "en") for row in read_jsonl(iter005 / "filtered/query_candidates_keep.jsonl") if is_vstyle_en(row)]
    controls = [mark_language(row, "en") for row in read_jsonl(iter005 / "filtered/query_tts_qwen3_input_structured.jsonl") if is_vstyle_en(row)]
    eval_rows = [dict(row, language="en") for row in read_jsonl(iter005 / "eval_inputs/e2e_input_aligned.jsonl") if is_vstyle_en(row)]

    counts = {
        "vstyle_en_candidates": write_jsonl(out_dir / "query_candidates.jsonl", candidates),
        "vstyle_en_candidates_keep": write_jsonl(out_dir / "filtered/query_candidates_keep.jsonl", candidates),
        "vstyle_en_query_tts_controls": write_jsonl(out_dir / "filtered/query_tts_qwen3_input_structured.jsonl", controls),
        "vstyle_en_existing_eval_rows": write_jsonl(out_dir / "existing_eval_audio_rows.jsonl", eval_rows),
    }
    (out_dir / "README_iter006_vstyle_en.txt").write_text(
        "Use filtered/query_candidates_keep.jsonl as RESPONSE_IN for English response/thinker target regeneration. "
        "existing_eval_audio_rows.jsonl preserves iter005 query_audio_path for vstyle_en.\n",
        encoding="utf-8",
    )
    return {"out_dir": str(out_dir), "counts": counts}


def copy_zh_reuse(iter005: Path, iter006: Path) -> Dict[str, Any]:
    out_dir = iter006 / "zh_reuse_from_iter005"
    out_dir.mkdir(parents=True, exist_ok=True)
    specs = [
        ("query_candidates.jsonl", "query_candidates.jsonl"),
        ("filtered/query_candidates_keep.jsonl", "filtered/query_candidates_keep.jsonl"),
        ("filtered/query_tts_qwen3_input_structured.jsonl", "filtered/query_tts_qwen3_input_structured.jsonl"),
        ("thinker_targets/response_controls.jsonl", "thinker_targets/response_controls.jsonl"),
        ("thinker_targets/thinker_targets.jsonl", "thinker_targets/thinker_targets.jsonl"),
        ("eval_inputs/e2e_input_aligned.jsonl", "eval_inputs/e2e_input_aligned.jsonl"),
        ("eval_inputs/talker_input_all.jsonl", "eval_inputs/talker_input_all.jsonl"),
        ("eval_inputs/response_tts_input_aligned.jsonl", "eval_inputs/response_tts_input_aligned.jsonl"),
    ]
    counts: Dict[str, Any] = {}
    for rel_src, rel_dst in specs:
        rows = [row for row in read_jsonl(iter005 / rel_src) if is_zh_reuse(row)]
        counts[rel_dst] = write_jsonl(out_dir / rel_dst, rows)
    for rel in ["sft/thinker_emochange/train.jsonl", "sft/thinker_emochange/val.jsonl", "sft/thinker_emochange/manifest.json"]:
        counts[rel] = "not_copied_final_sft_rebuilt_from_merged_eval"
    return {"out_dir": str(out_dir), "counts": counts}


def merge_unique(parts: Iterable[Iterable[Dict[str, Any]]]) -> List[Dict[str, Any]]:
    out: Dict[str, Dict[str, Any]] = {}
    for rows in parts:
        for row in rows:
            cur = qid(row)
            if cur:
                out[cur] = row
    return [out[key] for key in sorted(out)]


def merge_eval_with_existing_audio(new_eval_rows: List[Dict[str, Any]], existing_audio_rows: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
    audio_by_qid = by_qid(existing_audio_rows)
    merged = []
    for row in new_eval_rows:
        cur = qid(row)
        audio = audio_by_qid.get(cur)
        if audio:
            row = dict(row)
            row["query_audio_path"] = audio.get("query_audio_path", row.get("query_audio_path", ""))
            row["query_audio_control_path"] = audio.get("query_audio_control_path", row.get("query_audio_control_path", ""))
        merged.append(row)
    return merged


def merge_final(iter005: Path, iter006: Path, story_en_root: Path, vstyle_en_root: Path) -> Dict[str, Any]:
    zh_root = iter006 / "zh_reuse_from_iter005"
    final_eval = iter006 / "eval_inputs"
    final_targets = iter006 / "thinker_targets"
    final_filtered = iter006 / "filtered"

    if not zh_root.exists():
        copy_zh_reuse(iter005, iter006)

    zh_candidates = read_jsonl(zh_root / "query_candidates.jsonl")
    zh_keep = read_jsonl(zh_root / "filtered/query_candidates_keep.jsonl")
    zh_controls = read_jsonl(zh_root / "filtered/query_tts_qwen3_input_structured.jsonl")
    zh_responses = read_jsonl(zh_root / "thinker_targets/response_controls.jsonl")
    zh_targets = read_jsonl(zh_root / "thinker_targets/thinker_targets.jsonl")
    zh_e2e = read_jsonl(zh_root / "eval_inputs/e2e_input_aligned.jsonl")
    zh_talker = read_jsonl(zh_root / "eval_inputs/talker_input_all.jsonl")
    zh_response_aligned = read_jsonl(zh_root / "eval_inputs/response_tts_input_aligned.jsonl")

    story_candidates = read_jsonl(story_en_root / "query_candidates.jsonl")
    story_keep = read_jsonl(story_en_root / "filtered/query_candidates_keep.jsonl")
    story_controls = read_jsonl(story_en_root / "filtered/query_tts_qwen3_input_structured.jsonl")
    story_responses = read_jsonl(story_en_root / "thinker_targets/response_controls.jsonl")
    story_targets = read_jsonl(story_en_root / "thinker_targets/thinker_targets.jsonl")
    story_e2e = read_jsonl(story_en_root / "eval_inputs/e2e_input_aligned.jsonl")
    story_talker = read_jsonl(story_en_root / "eval_inputs/talker_input_all.jsonl")
    story_response_aligned = read_jsonl(story_en_root / "eval_inputs/response_tts_input_aligned.jsonl")

    vstyle_candidates = read_first_jsonl([vstyle_en_root / "filtered/query_candidates_keep.jsonl", vstyle_en_root / "query_candidates_keep.jsonl"])
    vstyle_controls = read_first_jsonl([vstyle_en_root / "filtered/query_tts_qwen3_input_structured.jsonl", vstyle_en_root / "query_tts_qwen3_input_structured.jsonl"])
    vstyle_responses = read_jsonl(vstyle_en_root / "thinker_targets/response_controls.jsonl")
    vstyle_targets = read_jsonl(vstyle_en_root / "thinker_targets/thinker_targets.jsonl")
    vstyle_export = read_jsonl(vstyle_en_root / "eval_inputs/e2e_input_aligned.jsonl")
    vstyle_existing_audio = read_jsonl(vstyle_en_root / "existing_eval_audio_rows.jsonl")
    if not vstyle_export and vstyle_responses and vstyle_targets and vstyle_existing_audio:
        response_by_qid = by_qid(vstyle_responses)
        audio_by_qid = by_qid(vstyle_existing_audio)
        vstyle_export = []
        for target in vstyle_targets:
            cur = qid(target)
            response = response_by_qid.get(cur, {})
            audio = audio_by_qid.get(cur, {})
            style = str((target.get("combined") or {}).get("instruct") or "").strip()
            text = str((target.get("combined") or {}).get("txt") or "").strip()
            visible_query = response.get("visible_query") or (response.get("source_query_candidate") or {}).get("visible_query") or target.get("source_query") or {}
            vstyle_export.append(
                {
                    "id": cur,
                    "qid": cur,
                    "query_type": target.get("query_type") or response.get("query_type"),
                    "language": "en",
                    "query_text": visible_query.get("text", audio.get("query_text", "")) if isinstance(visible_query, dict) else "",
                    "query_audio_path": audio.get("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"),
                }
            )
    vstyle_e2e = merge_eval_with_existing_audio(vstyle_export, vstyle_existing_audio)
    vstyle_talker = read_jsonl(vstyle_en_root / "eval_inputs/talker_input_all.jsonl")
    vstyle_response_aligned = read_jsonl(vstyle_en_root / "eval_inputs/response_tts_input_aligned.jsonl")
    if not vstyle_talker and vstyle_targets:
        vstyle_talker = []
        for target in vstyle_targets:
            cur = qid(target)
            combined = target.get("combined") or {}
            vstyle_talker.append(
                {
                    "id": cur,
                    "qid": cur,
                    "language": "en",
                    "ability": "constructed_thinker_to_talker",
                    "query_type": target.get("query_type"),
                    "thinker_style": str(combined.get("instruct") or "").strip(),
                    "thinker_text": str(combined.get("txt") or "").strip(),
                    "need_emochange": target.get("need_emochange"),
                    "source_query": target.get("source_query"),
                    "aligned_with_query_audio": cur in {qid(row) for row in vstyle_e2e},
                }
            )
    if not vstyle_response_aligned and vstyle_responses and vstyle_e2e:
        e2e_qids = {qid(row) for row in vstyle_e2e}
        vstyle_response_aligned = [row for row in vstyle_responses if qid(row) in e2e_qids]

    counts = {
        "query_candidates": write_jsonl(iter006 / "query_candidates.jsonl", merge_unique([zh_candidates, story_candidates, vstyle_candidates])),
        "query_candidates_keep": write_jsonl(final_filtered / "query_candidates_keep.jsonl", merge_unique([zh_keep, story_keep, vstyle_candidates])),
        "query_tts_controls": write_jsonl(final_filtered / "query_tts_qwen3_input_structured.jsonl", merge_unique([zh_controls, story_controls, vstyle_controls])),
        "response_controls": write_jsonl(final_targets / "response_controls.jsonl", merge_unique([zh_responses, story_responses, vstyle_responses])),
        "thinker_targets": write_jsonl(final_targets / "thinker_targets.jsonl", merge_unique([zh_targets, story_targets, vstyle_targets])),
        "e2e_input_aligned": write_jsonl(final_eval / "e2e_input_aligned.jsonl", merge_unique([zh_e2e, story_e2e, vstyle_e2e])),
        "talker_input_all": write_jsonl(final_eval / "talker_input_all.jsonl", merge_unique([zh_talker, story_talker, vstyle_talker])),
        "response_tts_input_aligned": write_jsonl(final_eval / "response_tts_input_aligned.jsonl", merge_unique([zh_response_aligned, story_response_aligned, vstyle_response_aligned])),
    }
    (iter006 / "iter006_merge_summary.json").write_text(json.dumps(counts, ensure_ascii=False, indent=2), encoding="utf-8")
    return {"iter006": str(iter006), "counts": counts}


def counts(iter006: Path) -> Dict[str, Any]:
    paths = [
        "query_candidates.jsonl",
        "filtered/query_candidates_keep.jsonl",
        "filtered/query_tts_qwen3_input_structured.jsonl",
        "thinker_targets/response_controls.jsonl",
        "thinker_targets/thinker_targets.jsonl",
        "eval_inputs/e2e_input_aligned.jsonl",
        "eval_inputs/talker_input_all.jsonl",
        "eval_inputs/response_tts_input_aligned.jsonl",
        "sft/thinker_emochange/train.jsonl",
        "sft/thinker_emochange/val.jsonl",
    ]
    return {rel: len(read_jsonl(iter006 / rel)) for rel in paths}


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("mode", choices=["copy-zh", "extract-vstyle-en", "merge-final", "counts"])
    parser.add_argument("--iter005", type=Path, default=DEFAULT_ITER005)
    parser.add_argument("--iter006", type=Path, default=DEFAULT_ITER006)
    parser.add_argument("--story_en_root", type=Path, default=DEFAULT_ITER006 / "english_rebuild/story_en")
    parser.add_argument("--vstyle_en_root", type=Path, default=DEFAULT_ITER006 / "english_rebuild/vstyle_en")
    args = parser.parse_args()

    if args.mode == "copy-zh":
        result = copy_zh_reuse(args.iter005, args.iter006)
    elif args.mode == "extract-vstyle-en":
        result = extract_vstyle_en(args.iter005, args.iter006)
    elif args.mode == "merge-final":
        result = merge_final(args.iter005, args.iter006, args.story_en_root, args.vstyle_en_root)
    else:
        result = {"iter006": str(args.iter006), "counts": counts(args.iter006)}
    print(json.dumps(result, ensure_ascii=False, indent=2))
    return 0


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