#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""Sample aligned reference-inference inputs by qid."""

from __future__ import annotations

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


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:
                rows.append(json.loads(line))
    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 row_id(row: Dict[str, Any]) -> str:
    return str(row.get("qid") or row.get("id") or "")


def filter_by_qids(rows: Iterable[Dict[str, Any]], qids: Set[str]) -> List[Dict[str, Any]]:
    return [row for row in rows if row_id(row) in qids]


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--e2e_input", type=Path, required=True)
    parser.add_argument("--talker_input", type=Path, required=True)
    parser.add_argument("--response_input", type=Path, required=True)
    parser.add_argument("--out_dir", type=Path, required=True)
    parser.add_argument("--sample_size", type=int, default=50)
    parser.add_argument("--seed", type=int, default=42)
    args = parser.parse_args()

    e2e_rows = read_jsonl(args.e2e_input)
    if args.sample_size > 0 and len(e2e_rows) > args.sample_size:
        e2e_rows = random.Random(args.seed).sample(e2e_rows, args.sample_size)
    e2e_rows = sorted(e2e_rows, key=row_id)
    qids = {row_id(row) for row in e2e_rows if row_id(row)}

    talker_rows = sorted(filter_by_qids(read_jsonl(args.talker_input), qids), key=row_id)
    response_rows = sorted(filter_by_qids(read_jsonl(args.response_input), qids), key=row_id)

    written = {
        "e2e_input_aligned": write_jsonl(args.out_dir / "e2e_input_aligned.jsonl", e2e_rows),
        "talker_input_all": write_jsonl(args.out_dir / "talker_input_all.jsonl", talker_rows),
        "response_tts_input_aligned": write_jsonl(args.out_dir / "response_tts_input_aligned.jsonl", response_rows),
    }
    summary = {
        "source": {
            "e2e_input": str(args.e2e_input),
            "talker_input": str(args.talker_input),
            "response_input": str(args.response_input),
        },
        "out_dir": str(args.out_dir),
        "seed": args.seed,
        "sample_size": args.sample_size,
        "selected_qids": sorted(qids),
        "written": written,
    }
    args.out_dir.mkdir(parents=True, exist_ok=True)
    (args.out_dir / "sample_summary.json").write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8")
    print(json.dumps(summary, ensure_ascii=False, indent=2))
    return 0


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