"""Parse release files + build the group index + enumerate training pairs.
FROZEN — do not modify (these define the task and the CV labels).
"""
import json
import os
from collections import defaultdict
from typing import Dict, List, Tuple

from .schema import Option, Question, TrainItem


def group_id(qid: str) -> str:
    return qid.rsplit("_", 1)[0]               # drops _1/_2/_3 or _e1..._e4


def build_index(cfg: dict) -> Tuple[List[Question], Dict[str, List[str]]]:
    qs: List[Question] = []
    groups: Dict[str, List[str]] = defaultdict(list)
    for name, rel in cfg["subsets"].items():
        path = os.path.join(cfg["data_root"], rel)
        base = os.path.dirname(path)
        for r in json.load(open(path, encoding="utf-8")):
            opts = [Option(letter=k[-1], wav=os.path.join(base, v))
                    for k, v in sorted(r["options"].items())]
            q = Question(
                subset=name, qid=r["question_id"], group=group_id(r["question_id"]),
                utterance_text=r["utterance"],
                utterance_wav=os.path.join(base, r["utterance_audio"]),
                context=r["context"], reference=r["response"], options=opts)
            qs.append(q)
            groups[q.group].append(q.qid)
    # Hard invariants double as a download sanity check.
    assert len(qs) == 530, f"expected 530 questions, got {len(qs)}"
    assert len(groups) == 200, f"expected 200 groups, got {len(groups)}"
    return qs, dict(groups)


def build_train_items(cfg: dict) -> List[TrainItem]:
    """Enumerate labelled (utterance, goodPara, badPara) triples for CV.

    Grouped by dialogue id so the group-aware CV split keeps a dialogue together,
    mirroring the test structure (group = dialogue, items = context/tone variants).
    """
    root = cfg["data_root"]
    items: List[TrainItem] = []

    # multi-context (t): 2 contexts per entry, one emotion.
    d = os.path.join(root, "empatheticDialogue_t_multi-context")
    mc = os.path.join(root, cfg["train"]["multi_context"])
    if os.path.exists(mc):
        for e in map(json.loads, open(mc, encoding="utf-8")):
            emo = e["emotion"]
            for i, ctx in enumerate(e["contexts"], start=1):
                utt = f'{d}/user_audio/{e["id"]}_{i}_{emo}.wav'
                good = f'{d}/response_audio/{e["id"]}_{i}_goodPara.wav'
                bad = f'{d}/response_audio/{e["id"]}_{i}_badPara.wav'
                if all(os.path.exists(p) for p in (utt, good, bad)):
                    items.append(TrainItem(e["id"], "multi_context", str(i),
                                           utt, good, bad, ctx["context"],
                                           ctx["response"], emo))

    # multi-emotion (n): 1 context, 2 emotions per entry.
    d = os.path.join(root, "empatheticDialogue_n_multi-emotion")
    me = os.path.join(root, cfg["train"]["multi_emotion"])
    if os.path.exists(me):
        for e in map(json.loads, open(me, encoding="utf-8")):
            ctx = e["contexts"]["Context"]
            emos = sorted({k[:-len("_response")] for k in e["contexts"]
                           if k.endswith("_response")})
            for em in emos:
                utt = f'{d}/user_audio/{e["id"]}_{em}.wav'
                good = f'{d}/response_audio/{e["id"]}_{em}_goodPara.wav'
                bad = f'{d}/response_audio/{e["id"]}_{em}_badPara.wav'
                if all(os.path.exists(p) for p in (utt, good, bad)):
                    items.append(TrainItem(e["id"], "multi_emotion", em, utt,
                                           good, bad, ctx,
                                           e["contexts"][f"{em}_response"], em))
    return items