#!/usr/bin/env python3
"""Build the Stage 1 SFT mix from downloaded raw splits + the manifest.

Pipeline per source key in ``training/configs/datasets.yaml``:

  raw.jsonl  --adapter-->  Example(s)  --think routing-->  ready / to_synthesize
                                       --dedup, source caps, data card

Outputs (default mix name ``stage1``):
  data/processed/stage1.ready.normalized.jsonl   rows with a <think> block, ready to split
  data/think_synthesis/stage1.to_synthesize.jsonl rows needing rejection-sampled reasoning
  reports/data/stage1_data_card.md               provenance + counts + license/cap flags

Typical flow:
  1. python training/scripts/hf_download.py --all --profile pilot
  2. python training/scripts/build_sft_dataset.py --profile pilot
  3. python training/scripts/synthesize_think.py --input data/think_synthesis/stage1.to_synthesize.jsonl ...
  4. python training/scripts/build_sft_dataset.py --include-synthesized data/think_synthesis/stage1.synthesized.jsonl
  5. python training/scripts/split_jsonl.py --input data/processed/stage1.ready.normalized.jsonl ...

This script is stdlib-only (plus PyYAML) so it runs on any host.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import random
import re
import sys
from pathlib import Path
from typing import Any, Iterable

import yaml

sys.path.insert(0, str(Path(__file__).resolve().parent))
from sft_adapters import Example, apply_adapter, has_think  # noqa: E402

WS_RE = re.compile(r"\s+")

WRAP_PLACEHOLDER = (
    "<think>\n"
    "Placeholder reasoning inserted for pipeline smoke testing only; replace "
    "with a synthesized trace before the real Stage 1 run.\n"
    "</think>\n\n"
)


def read_yaml(path: str | Path) -> dict[str, Any]:
    with Path(path).open("r", encoding="utf-8") as fh:
        payload = yaml.safe_load(fh) or {}
    if not isinstance(payload, dict):
        raise TypeError(f"Expected a YAML mapping in {path}")
    return payload


def read_jsonl(path: Path) -> Iterable[dict[str, Any]]:
    with path.open("r", encoding="utf-8") as fh:
        for line_no, line in enumerate(fh, start=1):
            line = line.strip()
            if not line:
                continue
            try:
                row = json.loads(line)
            except json.JSONDecodeError as exc:
                raise ValueError(f"Invalid JSON in {path}:{line_no}: {exc}") from exc
            if isinstance(row, dict):
                yield row


def estimate_tokens(messages: list[dict[str, str]]) -> int:
    chars = sum(len(m.get("content", "")) for m in messages)
    return max(1, chars // 4)


def dedup_key(messages: list[dict[str, str]]) -> str:
    text = " ".join(
        m.get("content", "") for m in messages if m.get("role") in {"user", "assistant"}
    )
    norm = WS_RE.sub(" ", text).strip().lower()
    return hashlib.sha256(norm.encode("utf-8")).hexdigest()


def assistant_turn_count(messages: list[dict[str, str]]) -> int:
    return sum(1 for m in messages if m.get("role") == "assistant")


def split_prompt_answer(messages: list[dict[str, str]]) -> tuple[list[dict[str, str]], str]:
    """Return (messages_without_final_assistant, final_assistant_content)."""
    for i in range(len(messages) - 1, -1, -1):
        if messages[i].get("role") == "assistant":
            return messages[:i], messages[i].get("content", "")
    return messages, ""


def wrap_messages(messages: list[dict[str, str]]) -> list[dict[str, str]]:
    out = []
    for m in messages:
        if m.get("role") == "assistant" and not has_think(m.get("content", "")):
            out.append({"role": "assistant", "content": WRAP_PLACEHOLDER + m.get("content", "")})
        else:
            out.append(m)
    return out


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--manifest", default="training/configs/datasets.yaml")
    parser.add_argument("--mix-name", default="stage1")
    parser.add_argument("--profile", choices=["full", "pilot"], default="full")
    parser.add_argument("--raw-dir", default=None, help="Override manifest defaults.raw_dir.")
    parser.add_argument("--processed-dir", default="data/processed")
    parser.add_argument("--synth-dir", default="data/think_synthesis")
    parser.add_argument("--report", default=None, help="Data card path (default reports/data/<mix>_data_card.md).")
    parser.add_argument(
        "--missing-think-policy",
        choices=["synthesize", "wrap", "drop"],
        default="synthesize",
    )
    parser.add_argument("--only", default=None, help="Comma-separated subset of source keys.")
    parser.add_argument("--include-disabled", action="store_true")
    parser.add_argument("--grounded-think", action="store_true",
                        help="Compose label-consistent <think> from metadata (no teacher pass needed).")
    parser.add_argument("--balance-detection", action="store_true",
                        help="Downsample majority-class vuln-detection rows to 1:1 so training doesn't collapse to majority.")
    parser.add_argument("--enforce-caps", action="store_true", help="Deterministically downsample over-cap sources.")
    parser.add_argument(
        "--include-synthesized",
        default=None,
        help="Fold an already-synthesized JSONL (from synthesize_think.py) into the ready mix.",
    )
    return parser.parse_args()


def cap_for(source: dict[str, Any], defaults: dict[str, Any], profile: str) -> int | None:
    if profile == "pilot":
        return source.get("pilot_sample_cap", defaults.get("pilot_sample_cap"))
    return source.get("sample_cap", defaults.get("sample_cap"))


def select_keys(sources: dict[str, Any], args: argparse.Namespace) -> list[str]:
    if args.only:
        wanted = [k.strip() for k in args.only.split(",") if k.strip()]
        for k in wanted:
            if k not in sources:
                raise SystemExit(f"Unknown source key in --only: {k}")
        return wanted
    return [k for k, s in sources.items() if s.get("enabled", False) or args.include_disabled]


def build() -> int:
    args = parse_args()
    manifest = read_yaml(args.manifest)
    defaults = manifest.get("defaults", {})
    raw_dir = Path(args.raw_dir or defaults.get("raw_dir", "data/download"))
    sources = manifest.get("sources", {})

    processed_dir = Path(args.processed_dir)
    synth_dir = Path(args.synth_dir)
    processed_dir.mkdir(parents=True, exist_ok=True)
    synth_dir.mkdir(parents=True, exist_ok=True)

    ready_path = processed_dir / f"{args.mix_name}.ready.normalized.jsonl"
    synth_path = synth_dir / f"{args.mix_name}.to_synthesize.jsonl"
    report_path = Path(args.report or f"reports/data/{args.mix_name}_data_card.md")
    report_path.parent.mkdir(parents=True, exist_ok=True)

    keys = select_keys(sources, args)

    seen: set[str] = set()
    per_source: dict[str, dict[str, Any]] = {}
    ready_rows: list[dict[str, Any]] = []
    synth_rows: list[dict[str, Any]] = []
    skipped: list[str] = []

    for key in keys:
        source = sources[key]
        raw_path = raw_dir / key / "raw.jsonl"
        stats = per_source.setdefault(
            key,
            {
                "hf_id": source.get("hf_id"),
                "group": source.get("group"),
                "license": source.get("license"),
                "auth": source.get("auth"),
                "adapter": source.get("adapter"),
                "raw_rows": 0,
                "examples": 0,
                "ready": 0,
                "to_synthesize": 0,
                "wrapped": 0,
                "dropped_no_think": 0,
                "dropped_dup": 0,
                "tokens": 0,
            },
        )
        if not raw_path.is_file():
            skipped.append(f"{key}: missing {raw_path} (run hf_download.py --key {key})")
            continue

        cap = cap_for(source, defaults, args.profile)
        adapter = source["adapter"]
        params = dict(source.get("params", {}) or {})
        if args.grounded_think:
            params["grounded_think"] = True
        kept_from_source = 0

        for row in read_jsonl(raw_path):
            stats["raw_rows"] += 1
            if cap is not None and kept_from_source >= cap:
                break
            try:
                examples = apply_adapter(adapter, row, params)
            except Exception as exc:  # noqa: BLE001 - one bad row shouldn't kill the build
                skipped.append(f"{key}: adapter error on row {stats['raw_rows']}: {exc!r}")
                continue
            for ex in examples:
                stats["examples"] += 1
                key_hash = dedup_key(ex.messages)
                if key_hash in seen:
                    stats["dropped_dup"] += 1
                    continue
                seen.add(key_hash)

                row_id = f"{source['hf_id']}:{key_hash[:16]}"
                tokens = estimate_tokens(ex.messages)

                if ex.think_status == "present" or all(
                    has_think(m["content"]) for m in ex.messages if m["role"] == "assistant"
                ):
                    ready_rows.append(_wrap_record(row_id, source, ex, ex.messages))
                    stats["ready"] += 1
                    stats["tokens"] += tokens
                    kept_from_source += 1
                    continue

                # needs reasoning synthesis
                if args.missing_think_policy == "drop":
                    stats["dropped_no_think"] += 1
                    continue
                if args.missing_think_policy == "wrap":
                    ready_rows.append(_wrap_record(row_id, source, ex, wrap_messages(ex.messages)))
                    stats["wrapped"] += 1
                    stats["ready"] += 1
                    stats["tokens"] += tokens
                    kept_from_source += 1
                    continue

                # synthesize: only single-assistant-turn examples are eligible
                if assistant_turn_count(ex.messages) != 1:
                    stats["dropped_no_think"] += 1
                    continue
                prompt_messages, answer = split_prompt_answer(ex.messages)
                synth_rows.append(
                    {
                        "id": row_id,
                        "source": source["hf_id"],
                        "license": source.get("license", "missing"),
                        "group": source.get("group"),
                        "prompt_messages": prompt_messages,
                        "reference_answer": answer,
                        "verify": ex.verify or {"mode": "backfill", "answer": answer},
                        "metadata": ex.metadata,
                    }
                )
                stats["to_synthesize"] += 1
                kept_from_source += 1

    # Optionally fold in already-synthesized rows.
    synthesized_added = 0
    if args.include_synthesized:
        synth_in = Path(args.include_synthesized)
        if not synth_in.is_file():
            skipped.append(f"--include-synthesized: missing {synth_in}")
        else:
            for row in read_jsonl(synth_in):
                messages = row.get("messages")
                if not isinstance(messages, list):
                    continue
                if not any(m.get("role") == "assistant" and has_think(m.get("content", "")) for m in messages):
                    continue
                key_hash = dedup_key(messages)
                if key_hash in seen:
                    continue
                seen.add(key_hash)
                ready_rows.append(row)
                synthesized_added += 1

    bal_log: list[str] = []
    if args.balance_detection:
        ready_rows, bal = _balance_detection(ready_rows)
        bal_log.append(
            f"detection balanced: vulnerable={bal['det_pos']} not_vulnerable={bal['det_neg']} "
            f"-> kept {bal['kept_each']} each; {bal['other']} non-detection rows untouched"
        )

    if args.enforce_caps:
        ready_rows, cap_log = _enforce_source_caps(ready_rows, manifest, per_source)
    else:
        cap_log = []
    cap_log = bal_log + cap_log

    _write_jsonl(ready_path, ready_rows)
    _write_jsonl(synth_path, synth_rows)
    _write_data_card(
        report_path, args, manifest, per_source, ready_rows, synth_rows, skipped, cap_log, synthesized_added
    )

    summary = {
        "mix": args.mix_name,
        "profile": args.profile,
        "ready_rows": len(ready_rows),
        "to_synthesize_rows": len(synth_rows),
        "synthesized_added": synthesized_added,
        "ready_path": str(ready_path),
        "synth_path": str(synth_path),
        "data_card": str(report_path),
        "skipped": len(skipped),
    }
    print(json.dumps(summary, indent=2))
    return 0


def _wrap_record(row_id: str, source: dict[str, Any], ex: Example, messages: list[dict[str, str]]) -> dict[str, Any]:
    return {
        "id": row_id,
        "source": source["hf_id"],
        "license": source.get("license", "missing"),
        "group": source.get("group"),
        "messages": messages,
        "metadata": {**ex.metadata, "think_status": ex.think_status},
    }


def _write_jsonl(path: Path, rows: list[dict[str, Any]]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as out:
        for row in rows:
            out.write(json.dumps(row, ensure_ascii=False, sort_keys=True) + "\n")


def _enforce_source_caps(
    ready_rows: list[dict[str, Any]],
    manifest: dict[str, Any],
    per_source: dict[str, Any],
) -> tuple[list[dict[str, Any]], list[str]]:
    """Deterministically downsample any source above the token-fraction cap."""
    data_mix = manifest.get("source_caps", {})
    max_frac = float(data_mix.get("max_single_source_token_fraction", 0.40))
    by_source: dict[str, list[dict[str, Any]]] = {}
    for row in ready_rows:
        by_source.setdefault(row.get("source", "?"), []).append(row)

    total_tokens = sum(estimate_tokens(r["messages"]) for r in ready_rows)
    cap_tokens = int(max_frac * total_tokens) if total_tokens else 0
    log: list[str] = []
    kept: list[dict[str, Any]] = []
    for src, rows in by_source.items():
        rows_sorted = sorted(rows, key=lambda r: r.get("id", ""))
        running = 0
        src_kept = []
        for r in rows_sorted:
            t = estimate_tokens(r["messages"])
            # Always keep at least one row so capping never zeroes a source; a
            # single row larger than the cap is kept rather than dropped whole.
            if cap_tokens and src_kept and running + t > cap_tokens:
                continue
            running += t
            src_kept.append(r)
        if len(src_kept) < len(rows):
            log.append(f"{src}: capped {len(rows)} -> {len(src_kept)} rows (~{max_frac:.0%} token cap)")
        kept.extend(src_kept)
    return kept, log


def _balance_detection(rows: list[dict[str, Any]], seed: int = 1337) -> tuple[list[dict[str, Any]], dict[str, int]]:
    """1:1 downsample vuln-detection rows by label; leave all other rows untouched."""
    det_pos, det_neg, other = [], [], []
    for r in rows:
        m = r.get("metadata", {}) or {}
        if m.get("task") == "vuln_detection" and m.get("label") in ("vulnerable", "not_vulnerable"):
            (det_pos if m["label"] == "vulnerable" else det_neg).append(r)
        else:
            other.append(r)
    rng = random.Random(seed)
    rng.shuffle(det_pos)
    rng.shuffle(det_neg)
    n = min(len(det_pos), len(det_neg))
    combined = other + det_pos[:n] + det_neg[:n]
    rng.shuffle(combined)
    return combined, {"det_pos": len(det_pos), "det_neg": len(det_neg), "kept_each": n, "other": len(other)}


def _token_fractions(ready_rows: list[dict[str, Any]]) -> dict[str, float]:
    by_source: dict[str, int] = {}
    for row in ready_rows:
        by_source[row.get("source", "?")] = by_source.get(row.get("source", "?"), 0) + estimate_tokens(
            row["messages"]
        )
    total = sum(by_source.values()) or 1
    return {k: v / total for k, v in sorted(by_source.items(), key=lambda kv: -kv[1])}


def _write_data_card(
    path: Path,
    args: argparse.Namespace,
    manifest: dict[str, Any],
    per_source: dict[str, Any],
    ready_rows: list[dict[str, Any]],
    synth_rows: list[dict[str, Any]],
    skipped: list[str],
    cap_log: list[str],
    synthesized_added: int,
) -> None:
    fractions = _token_fractions(ready_rows)
    lines: list[str] = []
    lines.append(f"# Data Card — {args.mix_name} ({args.profile})")
    lines.append("")
    lines.append(f"- Manifest: `{args.manifest}`")
    lines.append(f"- Missing-think policy: `{args.missing_think_policy}`")
    lines.append(f"- Ready rows (have `<think>`): **{len(ready_rows)}**")
    lines.append(f"- Rows queued for reasoning synthesis: **{len(synth_rows)}**")
    lines.append(f"- Synthesized rows folded in this build: **{synthesized_added}**")
    lines.append("")
    lines.append("## Per-source")
    lines.append("")
    lines.append("| key | hf_id | adapter | license | auth | raw | examples | ready | to_synth | dup | tokens≈ |")
    lines.append("|---|---|---|---|---|---|---|---|---|---|---|")
    for key, s in per_source.items():
        lines.append(
            f"| {key} | {s['hf_id']} | {s['adapter']} | {s['license']} | {s.get('auth')} | "
            f"{s['raw_rows']} | {s['examples']} | {s['ready']} | {s['to_synthesize']} | "
            f"{s['dropped_dup']} | {s['tokens']} |"
        )
    lines.append("")
    lines.append("## Ready-mix token fraction by source")
    lines.append("")
    cap = manifest.get("source_caps", {}).get("max_single_source_token_fraction", 0.40)
    for src, frac in fractions.items():
        flag = "  ⚠️ over cap" if frac > float(cap) else ""
        lines.append(f"- {src}: {frac:.1%}{flag}")
    lines.append("")
    lines.append(f"Single-source token cap: {float(cap):.0%}")
    if cap_log:
        lines.append("")
        lines.append("## Cap enforcement")
        for entry in cap_log:
            lines.append(f"- {entry}")
    # License flags
    missing_lic = sorted({s["hf_id"] for s in per_source.values() if str(s["license"]).lower() == "missing"})
    if missing_lic:
        lines.append("")
        lines.append("## ⚠️ Sources with no stated license (record provenance / academic-use only)")
        for hf_id in missing_lic:
            lines.append(f"- {hf_id}")
    if skipped:
        lines.append("")
        lines.append("## Skipped / warnings")
        for entry in skipped[:200]:
            lines.append(f"- {entry}")
        if len(skipped) > 200:
            lines.append(f"- ... and {len(skipped) - 200} more")
    lines.append("")
    lines.append("> Decontamination against eval splits is a separate required step before training.")
    path.write_text("\n".join(lines) + "\n", encoding="utf-8")


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