datasets.py

"""
NeuroJenML — Dataset builder, merger, validator, and HF export.

Architecture:
  - Neon stores manifests (metadata, IDs, usage tracking)
  - HF stores the actual JSONL files (single source of truth)
  - Training pulls JSONL from HF by path, same file that's in the dataset viewer
  - No JSONL stored inline in Neon — keeps the DB lean and avoids drift

Flow: build → push to HF → store manifest → (merge) → validate → train → mark used
"""

import os
import io
import json
import time
import asyncio
from datetime import datetime

import store


# ─── JSONL builders ──────────────────────────────────────────────

# Maximum training examples generated from a single paper.
# Each map edge produces 3 examples; without a cap, papers with many edges
# (e.g. cardiovascular / MI papers with 12+ edges → 36+ examples) numerically
# dominate the training set and cause domain drift.
_MAX_EXAMPLES_PER_PAPER = 15


def _systemic_examples(extractions: list) -> list:
    examples = []
    for ext in extractions:
        fields = ext.get("fields", {})
        cls = ext.get("classification", {})
        tags = cls.get("systemic_tags", [])
        hypothesis = fields.get("hypothesis", "")
        results = fields.get("results", "")
        maps = fields.get("maps", []) or []

        # Track how many examples this single paper contributes so that papers
        # with many edges (e.g. cardiovascular papers with 12+ edges) cannot
        # dominate the training set.
        paper_example_count = 0

        for m in maps:
            if not (m.get("from") and m.get("to")):
                continue
            if paper_example_count >= _MAX_EXAMPLES_PER_PAPER:
                break
            relation = m.get("relation", "relates_to")

            # Build a rationale that is explicitly conditioned on THIS edge
            # (from → relation → to), not a raw slice of the paper body.
            # A short context snippet is appended for grounding; the edge
            # description is always the lead sentence so the model learns the
            # correct association rather than surface co-occurrence.
            context_snippet = (results or hypothesis)[:200].strip()
            rationale = (
                f"{m['from']} drives '{relation}', leading to {m['to']}."
                + (f" Context: {context_snippet}" if context_snippet else "")
            )

            # Task 1: Mechanism inference — given correlation, predict mechanism
            examples.append({
                "task_type": "mechanism_inference",
                "instruction": (
                    "Given an observed association between a peripheral/systemic factor and a "
                    "central (brain) outcome in Alzheimer's disease, state the most likely "
                    "bridging mechanism. Respond as JSON with keys 'mechanism', 'rationale'."
                ),
                "input": f"Association: '{m['from']}' is linked to '{m['to']}'. What mechanism bridges them?",
                "output": json.dumps({
                    "mechanism": relation,
                    "rationale": rationale,
                }),
                "systemic_tags": tags,
                "paper_id": ext.get("paper_id"),
                "provenance": ext.get("provenance", {}),
            })

            # Task 2: Forward prediction — given peripheral factor, predict central outcome
            examples.append({
                "task_type": "forward_prediction",
                "instruction": (
                    "You are a systemic Alzheimer's reasoning model. Given a peripheral "
                    "(body-level) finding, predict the downstream brain pathology and the "
                    "biological mechanism that links them. Respond as JSON with keys "
                    "'mechanism', 'central', 'rationale'."
                ),
                "input": f"Peripheral finding: {m['from']}.",
                "output": json.dumps({
                    "mechanism": relation,
                    "central": m["to"],
                    "rationale": rationale,
                }),
                "systemic_tags": tags,
                "paper_id": ext.get("paper_id"),
                "provenance": ext.get("provenance", {}),
            })

            # Task 3: Reverse prediction — given central outcome, predict peripheral cause
            examples.append({
                "task_type": "reverse_prediction",
                "instruction": (
                    "Given a central brain pathology in Alzheimer's disease, identify "
                    "the peripheral (body-level) factors that may contribute to it. "
                    "Respond as JSON with keys 'peripheral', 'mechanism', 'rationale'."
                ),
                "input": f"Central pathology: {m['to']}.",
                "output": json.dumps({
                    "peripheral": m["from"],
                    "mechanism": relation,
                    "rationale": rationale,
                }),
                "systemic_tags": tags,
                "paper_id": ext.get("paper_id"),
                "provenance": ext.get("provenance", {}),
            })

            paper_example_count += 3  # 3 tasks per edge

    return examples


def _reasoning_examples(extractions: list) -> list:
    """Generate chain-of-thought (CoT) reasoning training examples.

    Three reasoning task types that force the model to think rather than retrieve:

    reasoning_chain
        A 5-step structured CoT trace wrapped in <think>...</think> before the
        final answer.  Each step is labelled so the model learns to scaffold
        its own reasoning process.

    counterfactual_reasoning
        "If X were absent or successfully treated, what would happen to Y?"
        Forces causal-direction reasoning rather than rote association recall.

    multi_hop_chain
        When two edges from the same paper share a middle node (A->B, B->C),
        generate a chain question: "given only A, trace through to C."  This
        teaches the model to compose two-step inference rather than looking up
        a direct link.
    """
    examples = []
    for ext in extractions:
        fields = ext.get("fields", {})
        cls = ext.get("classification", {})
        tags = cls.get("systemic_tags", [])
        hypothesis = fields.get("hypothesis", "")
        results = fields.get("results", "")
        maps = fields.get("maps", []) or []
        paper_id = ext.get("paper_id")
        provenance = ext.get("provenance", {})

        paper_reasoning_count = 0

        # Index edges by their 'from' node so we can find chains A->B, B->C.
        edge_index: dict = {}
        for m in maps:
            if m.get("from") and m.get("to"):
                edge_index.setdefault(m["from"], []).append(
                    (m.get("relation", "relates_to"), m["to"])
                )

        for m in maps:
            if not (m.get("from") and m.get("to")):
                continue
            if paper_reasoning_count >= _MAX_EXAMPLES_PER_PAPER:
                break

            src = m["from"]
            tgt = m["to"]
            relation = m.get("relation", "relates_to")
            context_snippet = (results or hypothesis)[:200].strip()

            # ── Task A: reasoning_chain ──────────────────────────────────────
            # A labelled 5-step CoT inside <think> blocks, followed by the final
            # JSON answer.  The model learns to narrate each inference step.
            think_block = (
                "<think>\n"
                f"Step 1 - Identify domain: '{src}' is a peripheral/systemic factor; "
                f"'{tgt}' is a central brain outcome in Alzheimer's disease.\n"
                f"Step 2 - Recall known link: prior evidence associates {src} with {tgt} "
                f"through the '{relation}' pathway.\n"
                f"Step 3 - Biological mechanism: {src} drives '{relation}', which in the "
                f"CNS context results in {tgt}. "
                + (f"Supporting context: {context_snippet}." if context_snippet else "") + "\n"
                f"Step 4 - Confidence assessment: the relationship is supported by "
                f"experimental evidence from the source paper.\n"
                f"Step 5 - Final answer: the mechanism is '{relation}'; the downstream "
                f"central pathology is '{tgt}'.\n"
                "</think>"
            )
            examples.append({
                "task_type": "reasoning_chain",
                "instruction": (
                    "Think step-by-step before answering. "
                    "Given a peripheral/systemic factor, reason through the biological "
                    "pathway to identify the central (brain) outcome in Alzheimer's disease. "
                    "First show your reasoning inside <think>...</think> tags, then respond "
                    "with JSON: {\"mechanism\": ..., \"central\": ..., \"confidence\": ...}."
                ),
                "input": f"Peripheral factor: {src}.",
                "output": (
                    think_block + "\n"
                    + json.dumps({
                        "mechanism": relation,
                        "central": tgt,
                        "confidence": "supported by experimental evidence",
                    })
                ),
                "systemic_tags": tags,
                "paper_id": paper_id,
                "provenance": provenance,
            })
            paper_reasoning_count += 1

            # ── Task B: counterfactual_reasoning ─────────────────────────────
            # Force causal-direction reasoning: the model must reason about what
            # would *not* happen if the upstream factor were removed or treated.
            cf_think = (
                "<think>\n"
                f"Counterfactual premise: suppose '{src}' is absent or successfully treated.\n"
                f"Step 1 - Causal chain: '{src}' normally drives '{relation}', "
                f"which leads to '{tgt}'.\n"
                f"Step 2 - Counterfactual intervention: if '{src}' is removed, the "
                f"'{relation}' signal is attenuated or eliminated.\n"
                f"Step 3 - Downstream effect: without the '{relation}' driver, "
                f"the manifestation of '{tgt}' would be reduced or delayed.\n"
                f"Step 4 - Caveats: redundant pathways may partially compensate; "
                "the effect size depends on whether this is the primary upstream driver.\n"
                f"Step 5 - Conclusion: treating or removing '{src}' is predicted to "
                f"attenuate '{tgt}' via the '{relation}' mechanism.\n"
                "</think>"
            )
            examples.append({
                "task_type": "counterfactual_reasoning",
                "instruction": (
                    "Think step-by-step. "
                    "Given a counterfactual scenario: if a peripheral factor were absent "
                    "or successfully treated, reason through what would happen to the "
                    "downstream brain outcome. "
                    "Show your reasoning inside <think>...</think> tags, then respond "
                    "with JSON: {\"effect\": ..., \"mechanism\": ..., \"caveats\": ...}."
                ),
                "input": (
                    f"Counterfactual: if '{src}' were absent or successfully treated, "
                    f"what would happen to '{tgt}'?"
                ),
                "output": (
                    cf_think + "\n"
                    + json.dumps({
                        "effect": f"'{tgt}' would be reduced or delayed",
                        "mechanism": f"removal of '{src}' attenuates the '{relation}' pathway",
                        "caveats": "redundant upstream pathways may partially compensate",
                    })
                ),
                "systemic_tags": tags,
                "paper_id": paper_id,
                "provenance": provenance,
            })
            paper_reasoning_count += 1

            # ── Task C: multi_hop_chain ──────────────────────────────────────
            # Find a second edge B->C where B == tgt (current edge A->B).
            # Generate a two-hop chain question: "given only A, trace to C."
            for next_relation, next_tgt in edge_index.get(tgt, []):
                if paper_reasoning_count >= _MAX_EXAMPLES_PER_PAPER:
                    break
                hop_think = (
                    "<think>\n"
                    f"Multi-hop chain: '{src}' -> '{tgt}' -> '{next_tgt}'.\n"
                    f"Step 1 - First hop: '{src}' drives '{relation}', leading to '{tgt}'.\n"
                    f"Step 2 - Second hop: '{tgt}', now acting as a secondary driver, "
                    f"drives '{next_relation}', leading to '{next_tgt}'.\n"
                    f"Step 3 - Chain composition: the net effect of '{src}' on '{next_tgt}' "
                    f"is mediated sequentially through '{tgt}'.\n"
                    f"Step 4 - Direct vs. indirect: this is an indirect effect; "
                    f"there may be no direct single-step edge from '{src}' to '{next_tgt}'.\n"
                    f"Step 5 - Conclusion: '{src}' reaches '{next_tgt}' indirectly via "
                    f"'{tgt}', through '{relation}' then '{next_relation}'.\n"
                    "</think>"
                )
                examples.append({
                    "task_type": "multi_hop_chain",
                    "instruction": (
                        "Think step-by-step. "
                        "Given a starting peripheral factor, trace a two-step biological "
                        "chain to identify the final downstream brain outcome in "
                        "Alzheimer's disease. "
                        "Show your reasoning inside <think>...</think> tags, then respond "
                        "with JSON: {\"intermediate\": ..., \"final_outcome\": ..., "
                        "\"step1_mechanism\": ..., \"step2_mechanism\": ...}."
                    ),
                    "input": (
                        f"Starting factor: '{src}'. "
                        f"What is the downstream brain outcome after two biological steps?"
                    ),
                    "output": (
                        hop_think + "\n"
                        + json.dumps({
                            "intermediate": tgt,
                            "final_outcome": next_tgt,
                            "step1_mechanism": relation,
                            "step2_mechanism": next_relation,
                        })
                    ),
                    "systemic_tags": tags,
                    "paper_id": paper_id,
                    "provenance": provenance,
                })
                paper_reasoning_count += 1
                break  # one multi-hop per originating edge is sufficient

    return examples


def _therapeutic_reasoning_examples(extractions: list) -> list:
    """Generate therapeutic reasoning training examples.

    These force the model to think like a drug designer, not a data retriever.
    The model must reason about:
    - What drug targets could address a mechanism
    - What barriers exist to targeting a pathway
    - What side effects might arise from intervention
    - How to design experiments to test hypotheses
    """
    import random as _rand
    examples = []
    paper_count = {}

    for ext in extractions:
        fields = ext.get("fields", {})
        maps = fields.get("maps", []) or []
        hypothesis = fields.get("hypothesis", "")
        results = fields.get("results", "")
        paper_id = ext.get("paper_id", "unknown")

        if paper_count.get(paper_id, 0) >= _MAX_EXAMPLES_PER_PAPER:
            continue
        if not maps:
            continue

        for m in maps:
            src = m.get("from", "")
            tgt = m.get("to", "")
            relation = m.get("relation", "relates_to")
            if not src or not tgt:
                continue
            if paper_count.get(paper_id, 0) >= _MAX_EXAMPLES_PER_PAPER:
                break

            # Task: Therapeutic target identification
            think = (
                "<think>\n"
                f"The mechanism linking {src} to {tgt} in Alzheimer's disease "
                f"involves '{relation}'. To intervene therapeutically, I need to consider "
                f"where in this pathway a drug could act.\n\n"
                f"First, what is the biological nature of this pathway? The '{relation}' "
                f"mechanism connects a peripheral factor to a central brain outcome. "
                f"This means there are multiple potential intervention points: "
                f"reducing the peripheral driver, blocking the transport/signaling mechanism, "
                f"or protecting the brain from downstream effects.\n\n"
                f"What are the barriers? The blood-brain barrier limits direct CNS access. "
                f"Systemic interventions may have off-target effects. The timing of intervention "
                f"matters — early intervention before neurodegeneration may be more effective.\n\n"
                f"What drug modalities could work? Antibodies targeting the peripheral factor, "
                f"small molecules blocking the signaling pathway, or gene therapies modifying "
                f"the expression of key mediators.\n\n"
                f"The most promising approach depends on the specificity of the '{relation}' "
                f"mechanism and whether it is a primary driver or a compensatory pathway.\n"
                "</think>"
            )
            examples.append({
                "task_type": "therapeutic_reasoning",
                "instruction": (
                    "You are a neuropharmacologist designing a therapeutic strategy for "
                    "Alzheimer's disease. Analyze the given mechanism and reason through "
                    "potential drug targets, delivery challenges, and experimental approaches. "
                    "Show your reasoning in <think>...</think>, then respond with analysis."
                ),
                "input": (
                    f"Mechanism: {src} drives '{relation}', leading to {tgt}. "
                    f"What are the therapeutic implications and drug design considerations?"
                ),
                "output": (
                    think + "\n"
                    + f"The therapeutic strategy for targeting the '{relation}' pathway "
                    f"between {src} and {tgt} should focus on: (1) identifying the most "
                    f"druggable node in the pathway, (2) developing delivery mechanisms "
                    f"that cross the blood-brain barrier where needed, and "
                    f"(3) designing biomarkers to monitor treatment response. "
                    f"The key challenge is that systemic interventions targeting {src} "
                    f"may have unintended consequences on other physiological processes."
                ),
                "paper_id": paper_id,
            })
            paper_count[paper_id] = paper_count.get(paper_id, 0) + 1

            # Task: Experimental design reasoning
            think2 = (
                "<think>\n"
                f"To test whether {src} causally contributes to {tgt} via '{relation}', "
                f"I need to design an experiment that isolates this specific pathway.\n\n"
                f"A correlation study alone is insufficient — I need interventional evidence. "
                f"Options include: (1) a knockout model where the {src} pathway is disabled, "
                f"(2) an antibody blockade study reducing circulating {src}, or "
                f"(3) a longitudinal cohort tracking {src} levels and {tgt} progression.\n\n"
                f"Each approach has tradeoffs. Knockout models prove causality but may not "
                f"translate to human therapy. Antibody blockade is more translatable but "
                f"requires knowing the right timing window. Longitudinal studies establish "
                f"temporal relationships but cannot prove mechanism.\n\n"
                f"The strongest evidence would combine approaches: use the longitudinal data "
                f"to identify the critical time window, then use the interventional studies "
                f"to test causality within that window.\n"
                "</think>"
            )
            examples.append({
                "task_type": "experimental_design",
                "instruction": (
                    "Design an experimental approach to test whether a peripheral factor "
                    "causally contributes to brain pathology in Alzheimer's disease. "
                    "Reason through the experimental design in <think>...</think>, "
                    "then describe your recommended approach."
                ),
                "input": (
                    f"Hypothesis: {src} contributes to {tgt} via '{relation}'. "
                    f"Design an experiment to test this."
                ),
                "output": (
                    think2 + "\n"
                    + f"To test the causal link between {src} and {tgt}, I recommend "
                    f"a multi-pronged approach: first establish the temporal relationship "
                    f"through longitudinal biomarker tracking, then use targeted intervention "
                    f"(antibody blockade or pathway inhibition) to test whether reducing "
                    f"the {src} signal attenuates {tgt} progression. Controls should include "
                    f"pathway-specific inhibitors and sham-treated animals."
                ),
                "paper_id": paper_id,
            })
            paper_count[paper_id] = paper_count.get(paper_id, 0) + 1

    return examples


def _general_examples(records: list) -> list:
    examples = []
    for rec in records:
        params = rec.get("parameters", {})
        if not params:
            continue
        examples.append({
            "task_type": "general_context",
            "instruction": (
                "Given standard demographic/clinical parameters, assess Alzheimer's disease "
                "risk context. Respond as JSON with key 'assessment'."
            ),
            "input": json.dumps(params),
            "output": json.dumps({"assessment": "", "species_class": rec.get("species_class")}),
            "species_class": rec.get("species_class"),
            "dataset_name": rec.get("dataset_name"),
        })
    return examples


def _maps_examples(maps: list) -> list:
    examples = []
    for m in maps:
        if not (m.get("from") and m.get("to")):
            continue
        examples.append({
            "task_type": "map_forward",
            "instruction": "Predict the downstream node. Respond as JSON with key 'to'.",
            "input": json.dumps({"from": m["from"], "relation": m.get("relation")}),
            "output": json.dumps({"to": m["to"]}),
            "domain": m.get("domain", ""),
        })
        examples.append({
            "task_type": "map_reverse",
            "instruction": "Reverse-engineer the upstream node. Respond as JSON with key 'from'.",
            "input": json.dumps({"to": m["to"], "relation": m.get("relation")}),
            "output": json.dumps({"from": m["from"]}),
            "domain": m.get("domain", ""),
        })
    return examples


async def _gather_examples(kind: str) -> list:
    if kind == "systemic":
        extractions = await store.list_all("extractions", newest_first=True)
        return _systemic_examples(extractions)
    if kind == "general":
        records = await store.list_all("general_records", newest_first=True)
        return _general_examples(records)
    if kind == "maps":
        maps = await store.list_all("maps", newest_first=True)
        return _maps_examples(maps)
    if kind == "reasoning":
        # Reasoning examples are derived from the same extractions as systemic
        # examples.  They deliberately overlap in source data but teach a
        # completely different cognitive skill: structured inference rather than
        # pattern-matched retrieval.
        extractions = await store.list_all("extractions", newest_first=True)
        return _reasoning_examples(extractions)
    if kind == "therapeutic":
        # Therapeutic reasoning teaches the model to think like a drug designer:
        # analyzing mechanisms, identifying targets, designing experiments.
        extractions = await store.list_all("extractions", newest_first=True)
        return _therapeutic_reasoning_examples(extractions)
    if kind == "all_reasoning":
        # Combined reasoning + therapeutic examples for comprehensive training.
        extractions = await store.list_all("extractions", newest_first=True)
        return _reasoning_examples(extractions) + _therapeutic_reasoning_examples(extractions)
    return []


# ─── Normalization ───────────────────────────────────────────────

_CANONICAL_KEYS = ("task_type", "instruction", "input", "output",
                    "systemic_tags", "paper_id", "provenance",
                    "species_class", "dataset_name", "domain",
                    # reasoning examples carry their CoT inside output; no extra keys needed
                    )


def _normalize_example(ex: dict) -> dict:
    out = {}
    for k in _CANONICAL_KEYS:
        out[k] = ex.get(k, "" if k != "systemic_tags" else [])
    return out


def _to_jsonl(examples: list) -> str:
    return "\n".join(json.dumps(_normalize_example(ex), ensure_ascii=False) for ex in examples)


def parse_jsonl(jsonl: str) -> list:
    examples = []
    for line in jsonl.strip().split("\n"):
        line = line.strip()
        if not line:
            continue
        try:
            obj = json.loads(line)
            if isinstance(obj, dict):
                examples.append(obj)
        except (json.JSONDecodeError, ValueError):
            pass
    return examples


def merge_jsonl_parts(parts: list) -> str:
    """Merge multiple JSONL strings, deduplicating by (task_type, input) pair."""
    all_examples = []
    for part in parts:
        all_examples.extend(parse_jsonl(part))

    seen = set()
    unique = []
    for ex in all_examples:
        key = (ex.get("task_type", ""), ex.get("input", ""))
        if key not in seen:
            seen.add(key)
            unique.append(ex)

    return _to_jsonl(unique)


# ─── HF Hub operations ───────────────────────────────────────────

def _hf_token():
    return os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_TOKEN")


def _hf_repo():
    return os.getenv("HF_DATASET_REPO")


def _push_to_hf(kind: str, version: str, jsonl: str) -> dict:
    token = _hf_token()
    repo = _hf_repo()
    if not token or not repo:
        return {"pushed": False, "reason": "HF_TOKEN or HF_DATASET_REPO not set"}
    try:
        from huggingface_hub import HfApi
        api = HfApi(token=token)
        api.create_repo(repo, repo_type="dataset", exist_ok=True, private=True)
        path_in_repo = f"{kind}/{version}/training_data.jsonl"
        api.upload_file(
            path_or_fileobj=io.BytesIO(jsonl.encode("utf-8")),
            path_in_repo=path_in_repo,
            repo_id=repo,
            repo_type="dataset",
            commit_message=f"NeuroJenML {kind} dataset {version}",
        )
        return {"pushed": True, "repo": repo, "path": path_in_repo}
    except Exception as e:
        return {"pushed": False, "reason": str(e)[:300]}


def pull_from_hf(path_in_repo: str, max_retries: int = 3) -> str:
    """Pull a JSONL file from HF dataset repo by path. Retries on transient failures."""
    token = _hf_token()
    repo = _hf_repo()
    if not token or not repo:
        raise RuntimeError("HF_TOKEN or HF_DATASET_REPO not set")
    from huggingface_hub import hf_hub_download
    last_error = None
    for attempt in range(max_retries):
        try:
            local_path = hf_hub_download(
                repo_id=repo,
                filename=path_in_repo,
                repo_type="dataset",
                token=token,
            )
            with open(local_path, "r", encoding="utf-8") as f:
                return f.read()
        except Exception as e:
            last_error = e
            if attempt < max_retries - 1:
                import time as _time
                delay = 2 ** attempt
                print(f"[retry] HF pull attempt {attempt + 1}/{max_retries} failed: {str(e)[:100]}. Retrying in {delay}s...")
                _time.sleep(delay)
    raise last_error


def push_extraction_to_hf(paper_id: str, extraction: dict, repo: str, token: str) -> str:
    try:
        from huggingface_hub import HfApi
    except ImportError:
        raise RuntimeError("huggingface_hub not installed")
    import json as _json
    record = {
        "paper_id": paper_id,
        "model": extraction.get("model"),
        "classification": extraction.get("classification"),
        "fields": extraction.get("fields"),
        "provenance": extraction.get("provenance"),
        "extracted_at": datetime.now().isoformat(),
    }
    content = (_json.dumps(record, ensure_ascii=False) + "\n").encode("utf-8")
    path_in_repo = f"extractions/{paper_id}.jsonl"
    api = HfApi(token=token)
    api.upload_file(
        path_or_fileobj=content,
        path_in_repo=path_in_repo,
        repo_id=repo,
        repo_type="dataset",
        commit_message=f"add extraction {paper_id}",
    )
    return f"https://huggingface.co/datasets/{repo}/blob/main/{path_in_repo}"


# ─── Build & Store (manifest only, JSONL goes to HF) ─────────────

async def build_and_record(kind: str, push: bool = True) -> dict:
    """Build a dataset, push JSONL to HF, store manifest in Neon."""
    examples = await _gather_examples(kind)
    version = "v" + datetime.now().strftime("%Y%m%d-%H%M%S")
    jsonl = _to_jsonl(examples)

    push_meta = {"pushed": False}
    if push:
        push_meta = _push_to_hf(kind, version, jsonl)

    ds_id = f"ds-{kind}-{int(time.time() * 1000) % 10_000_000}"
    manifest = {
        "id": ds_id,
        "kind": kind,
        "version": version,
        "example_count": len(examples),
        "hf_path": push_meta.get("path", ""),
        "hf_repo": push_meta.get("repo", ""),
        "hf_pushed": push_meta.get("pushed", False),
        # Inline JSONL stored as fallback when HF push is unavailable.
        # This is what the training bridge uses when hf_path is empty.
        # Capped at 8 MB to stay within Neon JSONB limits.
        "jsonl_inline": jsonl if len(jsonl) < 8 * 1024 * 1024 else "",
        "source_papers": list({ex.get("paper_id") for ex in examples if ex.get("paper_id")}),
        "task_types": list({ex.get("task_type") for ex in examples if ex.get("task_type")}),
        "used_in_training": False,
        "training_job_id": None,
        "created_at": datetime.now().isoformat(),
    }
    await store.upsert("datasets", ds_id, manifest)

    return {
        "status": "ok",
        "dataset": manifest,
        "jsonl": jsonl,
        "count": len(examples),
    }


# ─── Dataset Operations ──────────────────────────────────────────

async def list_datasets(kind: str = None, unused_only: bool = False) -> list:
    filters = {}
    if kind:
        filters["kind"] = kind
    datasets = await store.query("datasets", filters, newest_first=True)
    if unused_only:
        datasets = [d for d in datasets if not d.get("used_in_training")]
    return datasets


async def get_dataset(ds_id: str) -> dict:
    return await store.get("datasets", ds_id) or {}


async def get_dataset_jsonl(ds_id: str) -> str:
    """Return the JSONL for a dataset, pulling from HF or falling back to inline.

    Priority:
      1. HF Hub (hf_path set and HF credentials available)
      2. Inline JSONL stored in the manifest at build time (jsonl_inline field)
    """
    ds = await store.get("datasets", ds_id)
    if not ds:
        return ""

    hf_path = ds.get("hf_path", "")
    if hf_path:
        try:
            return await asyncio.to_thread(pull_from_hf, hf_path)
        except Exception:
            pass  # fall through to inline

    return ds.get("jsonl_inline", "")


async def merge_datasets(ds_ids: list, name: str = None) -> dict:
    """Merge multiple datasets by pulling each from HF and combining."""
    all_examples = []
    source_ids = []
    source_papers = set()

    for ds_id in ds_ids:
        ds = await store.get("datasets", ds_id)
        if not ds:
            continue
        hf_path = ds.get("hf_path", "")
        if not hf_path:
            continue  # Skip datasets without HF path
        try:
            jsonl = await asyncio.to_thread(pull_from_hf, hf_path)
        except Exception:
            jsonl = ""
        if not jsonl:
            continue
        examples = parse_jsonl(jsonl)
        all_examples.extend(examples)
        source_ids.append(ds_id)
        source_papers.update(ds.get("source_papers", []))

    # Deduplicate by (task_type, input)
    seen = set()
    unique = []
    for ex in all_examples:
        key = (ex.get("task_type", ""), ex.get("input", ""))
        if key not in seen:
            seen.add(key)
            unique.append(ex)

    version = "v" + datetime.now().strftime("%Y%m%d-%H%M%S")
    jsonl = _to_jsonl(unique)

    # Push merged dataset to HF
    push_meta = _push_to_hf("merged", version, jsonl)

    ds_id = f"ds-merged-{int(time.time() * 1000) % 10_000_000}"
    manifest = {
        "id": ds_id,
        "kind": "merged",
        "version": version,
        "name": name or f"Merged from {len(source_ids)} datasets",
        "example_count": len(unique),
        "hf_path": push_meta.get("path", ""),
        "hf_repo": push_meta.get("repo", ""),
        "hf_pushed": push_meta.get("pushed", False),
        "source_datasets": source_ids,
        "source_papers": list(source_papers),
        "task_types": list({ex.get("task_type") for ex in unique if ex.get("task_type")}),
        "jsonl_inline": jsonl if len(jsonl) < 8 * 1024 * 1024 else "",
        "used_in_training": False,
        "training_job_id": None,
        "created_at": datetime.now().isoformat(),
    }
    await store.upsert("datasets", ds_id, manifest)

    return {
        "status": "ok",
        "dataset": manifest,
        "count": len(unique),
        "deduplicated": len(all_examples) - len(unique),
    }


async def validate_dataset(ds_id: str) -> dict:
    """Validate a dataset pulled from HF. Checks structure and content quality."""
    try:
        ds = await store.get("datasets", ds_id)
        jsonl = await asyncio.to_thread(
            pull_from_hf, (ds or {}).get("hf_path", "")
        )
    except Exception as e:
        return {"valid": False, "error": f"Could not pull from HF: {str(e)[:200]}"}

    if not jsonl or not jsonl.strip():
        return {"valid": False, "error": "Dataset is empty"}

    examples = parse_jsonl(jsonl)
    issues = []
    required_keys = {"task_type", "instruction", "input", "output"}
    seen_pairs = set()
    leakage_count = 0
    empty_fields = 0

    for i, ex in enumerate(examples):
        missing = required_keys - set(ex.keys())
        if missing:
            issues.append(f"Line {i + 1}: missing keys {missing}")
        for key in ["instruction", "input", "output"]:
            val = ex.get(key, "")
            if not val or (isinstance(val, str) and not val.strip()):
                empty_fields += 1
        pair = (ex.get("task_type", ""), ex.get("input", ""))
        if pair in seen_pairs:
            issues.append(f"Line {i + 1}: duplicate (task_type, input) pair")
        seen_pairs.add(pair)
        output_val = ex.get("output", "")
        input_val = ex.get("input", "")
        if isinstance(output_val, str) and isinstance(input_val, str):
            try:
                out_obj = json.loads(output_val)
                for v in out_obj.values():
                    if isinstance(v, str) and len(v) > 5 and v.lower() in input_val.lower():
                        leakage_count += 1
                        break
            except (json.JSONDecodeError, ValueError):
                pass

    total = len(examples)
    quality_score = 100.0
    if total == 0:
        quality_score = 0
    else:
        quality_score -= (empty_fields / max(1, total * 3)) * 30
        quality_score -= (leakage_count / total) * 40
        quality_score -= min(20, len(issues) * 2)
        quality_score = max(0, round(quality_score, 1))

    return {
        "valid": len(issues) == 0 and total > 0,
        "example_count": total,
        "quality_score": quality_score,
        "issues": issues[:20],
        "issue_count": len(issues),
        "empty_fields": empty_fields,
        "leakage_examples": leakage_count,
        "duplicate_pairs": len(examples) - len(seen_pairs),
    }


async def mark_used(ds_id: str, job_id: str) -> bool:
    ds = await store.get("datasets", ds_id)
    if not ds:
        return False
    ds["used_in_training"] = True
    ds["training_job_id"] = job_id
    ds["used_at"] = datetime.now().isoformat()
    await store.upsert("datasets", ds_id, ds)
    return True


async def mark_unused(ds_id: str) -> bool:
    """Allow a used dataset to be reused for training."""
    ds = await store.get("datasets", ds_id)
    if not ds:
        return False
    ds["used_in_training"] = False
    ds["reused"] = True
    ds["reused_at"] = datetime.now().isoformat()
    await store.upsert("datasets", ds_id, ds)
    return True