"""Build CORTEX code datasets from declared sources, with license traceability.

Every dataset CORTEX trains on must be declared in ``registry.json`` with its origin
and license. This script fetches a small, reproducible sample from each source and
writes two files per dataset:

  * ``datasets/<id>.jsonl``       the records
  * ``datasets/<id>.meta.json``   provenance: source, license, counts, SHA-256

It never touches the distributed model checkpoint.

Usage::

    python datasets/build.py --dataset cortex-code-sft
    python datasets/build.py --all --limit 200
"""

from __future__ import annotations

import argparse
import hashlib
import json
import urllib.parse
import urllib.request
from pathlib import Path

REPO_ROOT = Path(__file__).resolve().parent.parent
DATASETS_DIR = REPO_ROOT / "datasets"

__all__ = ["build_dataset", "load_registry"]


def load_registry(path: str | Path | None = None) -> dict:
    path = Path(path) if path else DATASETS_DIR / "registry.json"
    return json.loads(path.read_text(encoding="utf-8"))


def _rows(dataset: str, config: str, split: str, offset: int, length: int) -> dict:
    query = urllib.parse.urlencode(
        {"dataset": dataset, "config": config, "split": split,
         "offset": offset, "length": length}
    )
    req = urllib.request.Request(
        f"https://datasets-server.huggingface.co/rows?{query}",
        headers={"User-Agent": "cortex-dataset-builder"},
    )
    with urllib.request.urlopen(req, timeout=120) as r:
        return json.load(r)


def _fetch_all(dataset: str, config: str, split: str, limit: int, page: int = 100) -> list[dict]:
    out, offset = [], 0
    while len(out) < limit:
        payload = _rows(dataset, config, split, offset, min(page, limit - len(out)))
        batch = [row["row"] for row in payload.get("rows", [])]
        if not batch:
            break
        out.extend(batch)
        offset += len(batch)
    return out


def _sha256_text(text: str) -> str:
    return hashlib.sha256(text.encode("utf-8")).hexdigest()


def _record_from_source(entry: dict, row: dict) -> dict | None:
    """Map a source row to a CORTEX training record using the entry's field map."""
    fields = entry["fields"]
    if entry["record_type"] != "code-sft":
        return None
    prompt = row.get(fields["prompt"]) or ""
    completion = row.get(fields["completion"]) or ""
    if entry.get("instruction_field"):
        prompt = row.get(entry["instruction_field"]) or ""
        extra = row.get(entry.get("input_field", "")) or ""
        if extra:
            prompt = f"{prompt}\n{extra}".strip()
    if not prompt or not completion:
        return None
    source_id = row.get(entry["id_field"]) if entry.get("id_field") else None
    return {
        "text": f"### Instruction\n{prompt}\n\n### Response\n{completion}",
        "source_id": str(source_id) if source_id is not None else "",
    }


def build_dataset(entry: dict, limit: int) -> dict:
    print(f"[{entry['id']}] fetching {limit} rows from {entry['source']} ({entry['split']})")
    rows = _fetch_all(entry["source"], entry["config"], entry["split"], limit)

    records = [r for r in (_record_from_source(entry, row) for row in rows) if r is not None]
    if not records:
        raise ValueError(f"[{entry['id']}] produced no usable records")

    jsonl_path = DATASETS_DIR / f"{entry['id']}.jsonl"
    body = "\n".join(json.dumps(r, ensure_ascii=False) for r in records) + "\n"
    jsonl_path.write_text(body, encoding="utf-8")

    meta = {
        "id": entry["id"],
        "source": entry["source"],
        "config": entry["config"],
        "split": entry["split"],
        "origin": entry["origin"],
        "license": entry["license"],
        "license_url": entry["license_url"],
        "redistribution": entry["redistribution"],
        "attribution": entry.get("attribution"),
        "record_type": entry["record_type"],
        "rows_fetched": len(rows),
        "records_written": len(records),
        "jsonl_sha256": _sha256_text(body),
        "jsonl_bytes": len(body.encode("utf-8")),
        "note": entry.get("note"),
        "fetched_via": "https://datasets-server.huggingface.co/rows",
    }
    (DATASETS_DIR / f"{entry['id']}.meta.json").write_text(
        json.dumps(meta, indent=2, ensure_ascii=False) + "\n", encoding="utf-8"
    )
    print(f"[{entry['id']}] wrote {len(records)} records -> {jsonl_path.name} "
          f"({meta['jsonl_bytes']} bytes, sha256 {meta['jsonl_sha256'][:16]}...)")
    return meta


def main():
    parser = argparse.ArgumentParser(description="Build CORTEX datasets")
    parser.add_argument("--dataset", help="dataset id from registry.json")
    parser.add_argument("--all", action="store_true")
    parser.add_argument("--limit", type=int, default=200)
    args = parser.parse_args()

    entries = load_registry()["datasets"]
    if args.all:
        selected = entries
    elif args.dataset:
        selected = [e for e in entries if e["id"] == args.dataset]
        if not selected:
            raise SystemExit(f"unknown dataset id: {args.dataset}")
    else:
        raise SystemExit("pass --dataset <id> or --all")

    for entry in selected:
        build_dataset(entry, args.limit)


if __name__ == "__main__":
    main()