"""Shared text constants and raw-document readers."""
from __future__ import annotations

import glob
import gzip
import json
import os

import pyarrow.parquet as pq

DATA = os.environ.get("TA_DATA", "data")

# Special tokens get ids 0..len-1 in this order. Hermes-style tool tags so the later Hermes
# phase only has to change the system prompt and tool schema, not the token vocabulary.
SPECIAL_TOKENS = [
    "<|endoftext|>", "<|pad|>", "<|im_start|>", "<|im_end|>",
    "<think>", "</think>", "<tool_call>", "</tool_call>", "<tool_response>", "</tool_response>",
] + [f"<|reserved_{i}|>" for i in range(22)]
EOS_ID = 0
PAD_ID = 1

# stack-edu shard -> language (see scripts/download.py)
CODE_LANG = {}
for lang, shards in {"markdown": [0, 1], "python": [73, 74, 75, 76], "shell": [87, 88], "javascript": [50],
                     "typescript": [91], "go": [34], "rust": [86], "c": [18]}.items():
    for i in shards:
        CODE_LANG[i] = lang


def code_files(lang: str) -> list[str]:
    return [f"{DATA}/raw/code/stack-edu-{i:04d}.json.gz" for i, l in sorted(CODE_LANG.items()) if l == lang
            and os.path.exists(f"{DATA}/raw/code/stack-edu-{i:04d}.json.gz")]


def iter_code(path: str):
    with gzip.open(path, "rt") as f:
        for line in f:
            r = json.loads(line)
            p = (r.get("metadata") or {}).get("path") or ""
            yield (f"File: {p.lstrip('/')}\n\n" if p else "") + r["text"]


def iter_parquet(path: str):
    pf = pq.ParquetFile(path)
    for rg in range(pf.num_row_groups):
        for t in pf.read_row_group(rg, columns=["text"]).column("text").to_pylist():
            yield t


def source_files() -> dict[str, list[tuple[str, callable]]]:
    """Source name -> list of (file, iterator fn)."""
    out = {
        "english": [(f, iter_parquet) for f in sorted(glob.glob(f"{DATA}/raw/english/**/*.parquet", recursive=True))],
        "math": [(f, iter_parquet) for f in sorted(glob.glob(f"{DATA}/raw/math/**/*.parquet", recursive=True))],
    }
    for lang in sorted(set(CODE_LANG.values())):
        out[f"code_{lang}"] = [(f, iter_code) for f in code_files(lang)]
    return out