"""
Build the strict-format calibration slice and append it to the conversational
calibration set, producing the final 588-sample calibration file.

The conversational set (built by build_calib_set.py) is 100% chat-templated,
but constrained-output formats (tags, JSON, fixed labels) are served as raw
zero-shot prompts by many benchmark harnesses and applications. This script
adds 76 samples in RAW TEXT form using a zero-shot instruction scaffold
("### 指示/### 回答形式/### 入力:/### 応答:"), prompt + gold completion
concatenated (calibration is forward-only, so the model "sees" both the
constrained instruction and the correctly-formatted answer tokens):

  - 38 fixed-label NLI with <answer>-tags — content from JNLI *train*
    (zenless-lab/jnli parquet mirror; no evaluation-set overlap).
  - 22 short-answer QA with <answer>-tags — JA Wikipedia title → first
    sentence(s). Wiki stream skips 12000 articles to stay disjoint from
    held-out perplexity and the conversational set's knowledge slice.
  - 16 JSON-schema outputs — JA Wikipedia paragraph → strict JSON object.

Raw-text lines are stored as {"text": ...} (vs the conversational set's
{"messages": ...}); 03_ptq_modelopt.py encodes "text" lines directly without
the chat template.

Usage: python scripts/build_calib_strict_format.py
"""

import json
import random
import re
from pathlib import Path

from datasets import load_dataset

SEED = 1234
BASE_PATH = Path("data/calib/calib_conversational.jsonl")
OUT_PATH = Path("data/calib/calib_full.jsonl")
DATASHEET_PATH = Path("data/calib/calibration_datasheet_strict_format.json")
WIKI_SKIP = 12000

N_NLI = 38
N_QA = 22
N_JSON = 16

JNLI_LABEL = {"0": "yes", "1": "unknown", "2": "no", 0: "yes", 1: "unknown", 2: "no"}

NLI_TEMPLATE = """### 指示
前提と仮説の関係をyes、no、unknownの中から回答してください。

制約：
- 前提が仮説を含意する場合はyesと出力
- 前提が仮説の否定を含意する場合はnoと出力
- 前提が仮説を含意せず、その否定も含意しない場合はunknownと出力

### 回答形式
<answer></answer>タグで囲んで回答してください

### 入力:
前提：{premise}
仮説：{hypothesis}
### 応答:<answer>{label}</answer>"""

QA_TEMPLATE = """### 指示
質問に対して、簡潔に回答してください。

### 回答形式
<answer></answer>タグで囲んで回答してください

### 入力:
質問：「{title}」とは何ですか。
### 応答:<answer>{answer}</answer>"""

JSON_TEMPLATE = """### 指示
次の文章を読んで、記事の情報を指定されたJSONスキーマに従って出力してください。

### 回答形式
{{"title": string, "summary": string, "keywords": string[]}} のJSONオブジェクトのみを出力してください

### 入力:
{text}
### 応答:{json_out}"""


def first_sentences(text: str, max_chars: int) -> str:
    out = ""
    for sent in re.split(r"(?<=。)", text):
        if not sent.strip():
            continue
        if out and len(out) + len(sent) > max_chars:
            break
        out += sent
        if len(out) >= max_chars // 2:
            break
    return out.strip()


def top_keywords(text: str, title: str, k: int = 3):
    words = re.findall(r"[一-龥ァ-ヶー]{2,6}", text)
    freq = {}
    for w in words:
        if w in title:
            continue
        freq[w] = freq.get(w, 0) + 1
    return [w for w, _ in sorted(freq.items(), key=lambda kv: -kv[1])[:k]]


def build_nli(rng):
    ds = load_dataset("zenless-lab/jnli", split="train", streaming=True)
    per_label_target = {"yes": N_NLI - 2 * (N_NLI // 3), "no": N_NLI // 3, "unknown": N_NLI // 3}
    got = {"yes": [], "no": [], "unknown": []}
    for ex in ds:
        label = JNLI_LABEL.get(ex["label"])
        if label is None or len(got[label]) >= per_label_target[label]:
            if all(len(got[l]) >= per_label_target[l] for l in got):
                break
            continue
        got[label].append(
            {
                "text": NLI_TEMPLATE.format(
                    premise=ex["premise"].strip(), hypothesis=ex["hypothesis"].strip(), label=label
                ),
                "source": "zenless-lab/jnli#train",
                "kind": "strict_nli_answer_tag",
            }
        )
    samples = got["yes"] + got["no"] + got["unknown"]
    rng.shuffle(samples)
    return samples


def build_wiki_slices(rng):
    ds = load_dataset("wikimedia/wikipedia", "20231101.ja", split="train", streaming=True)
    it = iter(ds)
    for _ in range(WIKI_SKIP):
        next(it)
    qa, js = [], []
    for ex in it:
        text = ex["text"].strip()
        title = ex["title"].strip()
        if len(text) < 600 or len(title) > 30:
            continue
        if len(qa) < N_QA:
            ans = first_sentences(text, 220)
            if not ans or len(ans) < 20:
                continue
            qa.append(
                {
                    "text": QA_TEMPLATE.format(title=title, answer=ans),
                    "source": f"wikimedia/wikipedia:20231101.ja#{ex['id']}",
                    "kind": "strict_qa_answer_tag",
                }
            )
        elif len(js) < N_JSON:
            para = text[:1200]
            obj = {
                "title": title,
                "summary": first_sentences(text, 180),
                "keywords": top_keywords(para, title),
            }
            js.append(
                {
                    "text": JSON_TEMPLATE.format(
                        text=para, json_out=json.dumps(obj, ensure_ascii=False)
                    ),
                    "source": f"wikimedia/wikipedia:20231101.ja#{ex['id']}",
                    "kind": "strict_json_schema",
                }
            )
        else:
            break
    return qa, js


def main():
    rng = random.Random(SEED)
    base_lines = [json.loads(l) for l in BASE_PATH.read_text().splitlines() if l.strip()]
    assert len(base_lines) == 512, len(base_lines)

    print("Building strict-format NLI slice (JNLI train)...", flush=True)
    nli = build_nli(rng)
    print(f"  {len(nli)} NLI samples", flush=True)
    print("Building wiki QA + JSON slices...", flush=True)
    qa, js = build_wiki_slices(rng)
    print(f"  {len(qa)} QA, {len(js)} JSON samples", flush=True)
    strict = nli + qa + js
    assert len(strict) == N_NLI + N_QA + N_JSON, len(strict)

    all_lines = base_lines + strict
    rng.shuffle(all_lines)
    with OUT_PATH.open("w") as f:
        for s in all_lines:
            rec = {k: s[k] for k in ("messages", "text", "source", "kind") if k in s}
            f.write(json.dumps(rec, ensure_ascii=False) + "\n")

    datasheet = {
        "seed": SEED,
        "base": "calib_conversational.jsonl (512 samples, unchanged)",
        "total": len(all_lines),
        "strict_format_slice": {
            "count": len(strict),
            "share": round(len(strict) / len(all_lines), 4),
            "format": "raw text in llm-jp-eval's ### 指示/### 回答形式/### 入力/### 応答 scaffold "
                      "(NOT chat-templated at PTQ time)",
            "nli_answer_tag": {"count": len(nli), "source": "zenless-lab/jnli train split",
                                "license": "cc-by-sa-4.0 (JGLUE)"},
            "qa_answer_tag": {"count": len(qa), "source": "wikimedia/wikipedia 20231101.ja, "
                              f"skip {WIKI_SKIP} (disjoint from ppl held-out and the conversational set)",
                              "license": "cc-by-sa-4.0"},
            "json_schema": {"count": len(js), "source": "wikimedia/wikipedia 20231101.ja",
                             "license": "cc-by-sa-4.0"},
        },
        "rationale": "constrained-output formats are served as raw zero-shot prompts by many "
                     "harnesses and applications, so they are calibrated in that form as well.",
    }
    DATASHEET_PATH.write_text(json.dumps(datasheet, ensure_ascii=False, indent=2))
    print(f"Wrote {len(all_lines)} lines to {OUT_PATH}", flush=True)


if __name__ == "__main__":
    main()