import collections
import json
import os
import random
import re
from pathlib import Path
from huggingface_hub import hf_hub_download, list_repo_files
import pyarrow.parquet as pq

os.makedirs("/content/calib_parts", exist_ok=True)
if not Path("/content/calib_parts/ukisai/qwen38-agent-sft-combined-100pct.parquet").exists():
    hf_hub_download("ukisai/Qwen3.8-27B-multi-turn-agent-sft", "qwen38-agent-sft-combined-100pct.parquet", repo_type="dataset", local_dir="/content/calib_parts/ukisai")
if not Path("/content/calib_parts/whittle/data/train-00000-of-00051.parquet").exists():
    hf_hub_download("logic65/whittle-distill-qwen3.8-27b", "data/train-00000-of-00051.parquet", repo_type="dataset", local_dir="/content/calib_parts/whittle")
if not Path("/content/calib_parts/chitako/calibration_ja.txt").exists():
    hf_hub_download("ChiTako/japanese-imatrix-calibration", "calibration_ja.txt", repo_type="dataset", local_dir="/content/calib_parts/chitako")
if not Path("/content/calib_parts/hotchpotch/data/test-00000-of-00001.parquet").exists():
    hf_hub_download("hotchpotch/japanese-qa-reasoning-100k", "data/test-00000-of-00001.parquet", repo_type="dataset", local_dir="/content/calib_parts/hotchpotch")

rng = random.Random(20261014)

# 1. Load Ukisai and deduplicate by trial_name (keep final episode), then group-split by base task ID
tbl_uki = pq.read_table("/content/calib_parts/ukisai/qwen38-agent-sft-combined-100pct.parquet")
n_uki = tbl_uki.num_rows
tasks_col = tbl_uki["task"].to_pylist()
trials_col = tbl_uki["trial_name"].to_pylist()
episodes_col = tbl_uki["episode"].to_pylist()

def ep_num(ep_str):
    m = re.search(r"(\d+)$", str(ep_str))
    return int(m.group(1)) if m else 0

def get_base_task(t_str):
    s = str(t_str)
    s = re.sub(r"(__.*|_run\d+)$", "", s)
    return s

# Keep only the max episode index per trial_name
best_row_for_trial = {}
for idx in range(n_uki):
    tr = trials_col[idx]
    en = ep_num(episodes_col[idx])
    if tr not in best_row_for_trial or en > best_row_for_trial[tr][0]:
        best_row_for_trial[tr] = (en, idx)

# Group trials by base_task ( keep only 1 trial per base_task so zero base_task duplicates! )
task_to_row = {}
for tr, (_en, idx) in best_row_for_trial.items():
    bt = get_base_task(tasks_col[idx])
    if bt not in task_to_row:
        task_to_row[bt] = idx

nl2bash_tasks = sorted([bt for bt in task_to_row if bt.startswith("nl2bash")])
other_tasks   = sorted([bt for bt in task_to_row if not bt.startswith("nl2bash")])
rng.shuffle(nl2bash_tasks)
rng.shuffle(other_tasks)

# Disjoint split by base_task:
calib_uki_tasks = nl2bash_tasks[:240] + other_tasks[:60]
eval_uki_tasks  = nl2bash_tasks[240:275] + other_tasks[60:70]
assert set(calib_uki_tasks).isdisjoint(set(eval_uki_tasks))
rng.shuffle(calib_uki_tasks)
rng.shuffle(eval_uki_tasks)

def fmt_uki_row(row_idx, keep_preamble=False):
    conv = tbl_uki["conversations"][row_idx].as_py()
    parts = []
    for m_i, msg in enumerate(conv):
        role = msg.get("role", "user")
        content = str(msg.get("content", ""))
        if m_i == 0 and not keep_preamble and "Task Description:" in content:
            content = "Task Description:" + content.split("Task Description:", 1)[1]
        parts.append(f"<|im_start|>{role}\n{content}<|im_end|>")
    full = "\n".join(parts)
    # Keep tail if very long so multi-turn interactions and final turns are preserved
    if len(full) > 5000:
        full = full[:800] + "\n...\n" + full[-4200:]
    return full

uki_calib_docs = [fmt_uki_row(task_to_row[bt], keep_preamble=(i < 2)) for i, bt in enumerate(calib_uki_tasks)]
uki_eval_docs  = [fmt_uki_row(task_to_row[bt], keep_preamble=False) for bt in eval_uki_tasks]

# 2. Terminal-Bench 2.1: use xhigh/traces/<task>/prompt.txt, split strictly by <task>
files_tb = list_repo_files("Lottolabs/terminal-bench-2.1-qwen3.8-27b-traces", repo_type="dataset")
def get_tb_task_name(f):
    return re.sub(r"-[0-9a-f]{8,}$", "", f.split("/")[2])

tb_by_task = {}
for f in sorted(files_tb):
    if f.startswith("xhigh/traces/") and f.endswith("/prompt.txt"):
        tname = get_tb_task_name(f)
        if tname not in tb_by_task:
            tb_by_task[tname] = f

xhigh_prompts = list(tb_by_task.values())
rng.shuffle(xhigh_prompts)
tb_calib_files = xhigh_prompts[:24]
tb_eval_files  = xhigh_prompts[24:32]

assert set(map(get_tb_task_name, tb_calib_files)).isdisjoint(set(map(get_tb_task_name, tb_eval_files)))

def clean_tb_prompt(txt: str) -> str:
    if "\n\nExecution contract:" in txt:
        txt = txt.split("\n\nExecution contract:", 1)[0].strip()
    return txt[-4500:] if len(txt) > 4500 else txt

tb_calib_docs = []
for f in tb_calib_files:
    p = hf_hub_download("Lottolabs/terminal-bench-2.1-qwen3.8-27b-traces", f, repo_type="dataset", local_dir="/content/calib_parts/tb")
    txt = Path(p).read_text(encoding="utf-8", errors="replace")
    # Keep tb_calib_docs identical to Phase 1a so calibration_mix_v2.txt is 100% reproducible
    tb_calib_docs.append(txt[-4500:] if len(txt) > 4500 else txt)

tb_eval_docs = []
for f in tb_eval_files:
    p = hf_hub_download("Lottolabs/terminal-bench-2.1-qwen3.8-27b-traces", f, repo_type="dataset", local_dir="/content/calib_parts/tb")
    txt = Path(p).read_text(encoding="utf-8", errors="replace")
    tb_eval_docs.append(clean_tb_prompt(txt))

# 3. Whittle: use prompt + response directly (already templated!), stratify by kind, truncate prompt from left
tbl_whi = pq.read_table("/content/calib_parts/whittle/data/train-00000-of-00051.parquet")
whi_by_kind = collections.defaultdict(list)
for idx in range(tbl_whi.num_rows):
    k = str(tbl_whi["kind"][idx].as_py())
    whi_by_kind[k].append(idx)

for k in whi_by_kind:
    rng.shuffle(whi_by_kind[k])

def fmt_whi_row(idx):
    p = str(tbl_whi["prompt"][idx].as_py() or "")
    r = str(tbl_whi["response"][idx].as_py() or "")
    if not r.strip():
        reas = str(tbl_whi["reasoning"][idx].as_py() or "")
        ans  = str(tbl_whi["answer"][idx].as_py() or "")
        r = f"{reas}\n</think>\n{ans}<|im_end|>"
    if len(p) > 4200:
        # Keep system header + tail of prompt containing task & latest tool output
        p = p[:500] + "\n...\n" + p[-3700:]
    if len(r) > 3800:
        r = r[:1800] + "\n...\n" + r[-2000:]
    return p + r

# Stratified selection: all agentx (69), agent1 (39), json (47), longctx (64), chat (30), plus 24 math/gsm, 40 code
whi_calib_idxs = []
whi_eval_idxs  = []
quotas = [("agentx", 55, 12), ("agent1", 31, 8), ("json", 38, 9), ("longctx", 52, 12), ("chat", 24, 6), ("code", 35, 8), ("math", 15, 4), ("gsm", 10, 2)]
for k, n_cal, n_ev in quotas:
    pool = whi_by_kind.get(k, [])
    whi_calib_idxs.extend(pool[:n_cal])
    whi_eval_idxs.extend(pool[n_cal : n_cal + n_ev])

assert set(whi_calib_idxs).isdisjoint(set(whi_eval_idxs))
rng.shuffle(whi_calib_idxs)
rng.shuffle(whi_eval_idxs)
whi_calib_docs = [fmt_whi_row(i) for i in whi_calib_idxs]
whi_eval_docs  = [fmt_whi_row(i) for i in whi_eval_idxs]

# 4. Japanese:
#    - Calibration JA = deduplicated ChiTako + hotchpotch[:200] (by distinct url/id)
#    - Held-out Eval JA = hotchpotch[200:260] (strictly url/id disjoint, and 0 shingle overlap with ChiTako!)
ja_raw = Path("/content/calib_parts/chitako/calibration_ja.txt").read_text(encoding="utf-8", errors="replace")
raw_paras = [b.strip() for b in ja_raw.split("\n\n") if len(b.strip()) > 120]

def get_shingles(s, k=20):
    s_norm = re.sub(r"\s+", " ", s)
    if len(s_norm) < k:
        return {s_norm}
    return {s_norm[i : i + k] for i in range(0, len(s_norm) - k + 1, 5)}

dedup_chitako = []
seen_shingles = set()
for p in raw_paras:
    sh = get_shingles(p)
    if len(sh & seen_shingles) / max(1, len(sh)) > 0.30:
        continue
    seen_shingles |= sh
    dedup_chitako.append(p[:2000])
rng.shuffle(dedup_chitako)
print(f"ChiTako raw paras: {len(raw_paras)} -> deduplicated: {len(dedup_chitako)}")

tbl_hp = pq.read_table("/content/calib_parts/hotchpotch/data/test-00000-of-00001.parquet")
url_to_hp = {}
for idx in range(tbl_hp.num_rows):
    u = str(tbl_hp["url"][idx].as_py() or tbl_hp["id"][idx].as_py() or idx)
    if u not in url_to_hp:
        url_to_hp[u] = idx

hp_urls = sorted(url_to_hp.keys())
rng.shuffle(hp_urls)
hp_calib_docs = []
hp_eval_docs  = []
for u in hp_urls:
    idx = url_to_hp[u]
    q = str(tbl_hp["question"][idx].as_py() or "")
    r = str(tbl_hp["reasoning"][idx].as_py() or "")
    a = str(tbl_hp["answer"][idx].as_py() or "")
    ctx = str(tbl_hp["text"][idx].as_py() or "")[:1200]
    doc = f"<|im_start|>user\n以下の文章を参考に質問に答えてください。\n{ctx}\n\n質問: {q}<|im_end|>\n<|im_start|>assistant\n<think>\n{r[:1600]}\n</think>\n{a}<|im_end|>"
    sh = get_shingles(doc)
    if len(sh & seen_shingles) / max(1, len(sh)) > 0.25:
        continue
    seen_shingles |= sh
    if len(hp_calib_docs) < 180:
        hp_calib_docs.append(doc)
    elif len(hp_eval_docs) < 45:
        hp_eval_docs.append(doc)

# Group ChiTako paragraphs into ~1,600-char (~1,000-token) JA blocks
chitako_grouped = []
cur_g = []
cur_len = 0
for p in dedup_chitako:
    cur_g.append(p)
    cur_len += len(p)
    if cur_len >= 1500:
        chitako_grouped.append("\n\n".join(cur_g))
        cur_g = []
        cur_len = 0
if cur_g:
    chitako_grouped.append("\n\n".join(cur_g))

# Interleave round-robin at document granularity across all 5 sources, token-budgeted per round:
# In each round i:
#   - 1 Ukisai doc (~3,500 chars EN/bash ≈ 950 tokens)
#   - every 4th round: 1 Terminal-Bench xhigh prompt (~4,000 chars ≈ 1,100 tokens)
#   - 1 Whittle doc (~3,600 chars EN/code/json/think ≈ 1,000 tokens)
#   - 1 JA doc alternating ChiTako (~1,600 chars JA ≈ 1,000 tokens) and Hotchpotch (~1,600 chars JA ≈ 1,000 tokens)
ja_interleaved_docs = []
for i in range(max(len(chitako_grouped), len(hp_calib_docs))):
    if i < len(chitako_grouped): ja_interleaved_docs.append(chitako_grouped[i])
    if i < len(hp_calib_docs):   ja_interleaved_docs.append(hp_calib_docs[i])

n_rounds = max(len(uki_calib_docs), len(whi_calib_docs), len(ja_interleaved_docs))
rounds_text = []
for r in range(n_rounds):
    if r < len(uki_calib_docs):
        rounds_text.append(uki_calib_docs[r])
    if r % 4 == 0 and (r // 4) < len(tb_calib_docs):
        rounds_text.append(tb_calib_docs[r // 4])
    if r < len(whi_calib_docs):
        rounds_text.append(whi_calib_docs[r])
    if r < len(ja_interleaved_docs):
        rounds_text.append(ja_interleaved_docs[r])

calib_v2 = "\n\n".join(rounds_text)
Path("/content/calibration_mix_v2.txt").write_text(calib_v2, encoding="utf-8")

eval_ops_v2  = "\n\n".join(uki_eval_docs[:25] + tb_eval_docs)[:80_000]
eval_chat_v2 = "\n\n".join(whi_eval_docs[:30])[:80_000]
eval_ja_v2   = "\n\n".join(hp_eval_docs[:35])[:80_000]

Path("/content/eval_ops_tools.txt").write_text(eval_ops_v2, encoding="utf-8")
Path("/content/eval_chat_json.txt").write_text(eval_chat_v2, encoding="utf-8")
Path("/content/eval_ja.txt").write_text(eval_ja_v2, encoding="utf-8")

import hashlib
print(
    f"Per-source calib chars: ukisai={sum(len(x) for x in uki_calib_docs):,}, "
    f"tb={sum(len(x) for x in tb_calib_docs):,}, whittle={sum(len(x) for x in whi_calib_docs):,}, "
    f"chitako={sum(len(x) for x in chitako_grouped):,}, hotchpotch={sum(len(x) for x in hp_calib_docs):,}"
)
print(
    f"Per-source eval_ops chars before 80K cap: ukisai_25={sum(len(x) for x in uki_eval_docs[:25]):,}, "
    f"tb_8={sum(len(x) for x in tb_eval_docs):,}"
)
for p in ["/content/calibration_mix_v2.txt", "/content/eval_ops_tools.txt", "/content/eval_chat_json.txt", "/content/eval_ja.txt"]:
    raw_b = Path(p).read_bytes()
    h256 = hashlib.sha256(raw_b).hexdigest()
    print(f"Wrote {p:35s}: {len(raw_b) / 1024:8.1f} KB | sha256={h256}")