#!/usr/bin/env python3
"""Per-source schema adapters for the CyberGym SFT mix.

Each raw Hugging Face dataset in ``training/configs/datasets.yaml`` has its own
column layout. The generic ``normalize_sft_jsonl.py`` only handles a few common
shapes (``messages`` / ``conversations`` / ``instruction|input|output`` /
``system|user|assistant``) and silently rejects everything else. The Tier-1
C/C++ detection tables (PrimeVul, DiverseVul) and the vuln/fix pair sets
(CrossVul) do **not** match any of those shapes, so they need dedicated
adapters.

An adapter takes one raw row (a ``dict``) plus a small ``params`` dict from the
manifest and returns a list of :class:`Example` objects (0, 1, or 2 per row).
The orchestrator (``build_sft_dataset.py``) wraps each Example with provenance
(id/source/license) and routes it to the ``ready`` mix or the ``to_synthesize``
queue based on ``think_status``.

This module is deliberately stdlib-only so it can be unit-tested without the
``datasets``/``torch`` stack (those only exist on the rented GPU host).
"""

from __future__ import annotations

import json
from dataclasses import dataclass, field
from typing import Any, Callable

DEFAULT_SYSTEM = "Authorized security research and education context."

# A short, model-facing instruction reused by the code-analysis adapters so the
# synthesized <think> traces and answers stay in one consistent format.
DETECTION_SYSTEM = (
    "You are a senior security engineer reviewing source code for "
    "memory-safety and other security vulnerabilities. Reason carefully, then "
    "give a clear verdict. Authorized security research and education context."
)


@dataclass
class Example:
    """One normalized training example before provenance wrapping."""

    messages: list[dict[str, str]]
    # "present"  -> assistant turn already contains a <think> block
    # "needs_synthesis" -> reasoning must be backfilled before Stage 1
    think_status: str = "needs_synthesis"
    # Optional ground truth used by the rejection-sampling teacher pass.
    # {"mode": "label", "expected": "vulnerable", "cwe": [...]}  -> verify label
    # {"mode": "backfill", "answer": "..."}                       -> keep answer, add think
    verify: dict[str, Any] | None = None
    metadata: dict[str, Any] = field(default_factory=dict)


# --------------------------------------------------------------------------- #
# helpers
# --------------------------------------------------------------------------- #
def coerce_list(value: Any) -> list[str]:
    """Normalize a CWE-ish field into a clean list of strings.

    Handles native lists, JSON-encoded list strings (``'["CWE-310"]'`` as
    returned by the datasets-server), comma/space separated strings, and the
    various null sentinels seen in these datasets.
    """
    if value is None:
        return []
    if isinstance(value, list):
        items = value
    elif isinstance(value, str):
        s = value.strip()
        if not s or s.lower() in {"none", "null", "[]", "nan"}:
            return []
        if s.startswith("["):
            try:
                parsed = json.loads(s)
                items = parsed if isinstance(parsed, list) else [parsed]
            except json.JSONDecodeError:
                items = [p.strip() for p in s.strip("[]").split(",")]
        else:
            items = [p.strip() for p in s.replace(";", ",").split(",")]
    else:
        items = [value]
    out: list[str] = []
    for item in items:
        text = str(item).strip().strip("'\"")
        if text and text.lower() not in {"none", "null", "nan"}:
            out.append(text)
    return out


def as_bool_label(value: Any) -> bool | None:
    """Interpret a vulnerability label that may be 0/1, "0"/"1", or bool."""
    if value is None:
        return None
    if isinstance(value, bool):
        return value
    if isinstance(value, (int, float)):
        return int(value) == 1
    text = str(value).strip().lower()
    if text in {"1", "true", "vulnerable", "yes", "vuln"}:
        return True
    if text in {"0", "false", "not vulnerable", "no", "safe", "secure", "benign"}:
        return False
    return None


def get_first(row: dict[str, Any], keys: list[str]) -> Any:
    for key in keys:
        if key in row and row[key] not in (None, ""):
            return row[key]
    return None


def has_think(text: str) -> bool:
    return "<think>" in text and "</think>" in text


def ensure_think_closed(think: str) -> str:
    """Make sure a reasoning fragment is a well-formed <think> block."""
    t = think.strip()
    if not t:
        return ""
    if "<think>" not in t:
        t = "<think>\n" + t
    if "</think>" not in t:
        t = t + "\n</think>"
    return t


def fence(code: str, language: str) -> str:
    lang = (language or "").strip().lower()
    alias = {"c++": "cpp", "c/c++": "cpp", "cpp": "cpp", "c": "c"}.get(lang, lang or "")
    body = str(code).rstrip()
    return f"```{alias}\n{body}\n```"


def _compose_grounded(task: str, label: str, cwes: list[str], desc: str | None, language: str) -> str:
    """Compose a short, label-consistent <think> from ground-truth metadata.

    This is "grounded backfill": real reasoning derived from the dataset's own
    labels/descriptions, so detection data is training-ready with a <think> block
    without a slow teacher/rejection-sampling pass. The reasoning always agrees
    with the gold answer (never teaches a wrong conclusion).
    """
    cwe_txt = ", ".join(cwes) if cwes else ""
    lang = language or "code"
    if task == "secure_fix":
        cp = f" (it is affected by {cwe_txt})" if cwe_txt else ""
        return (
            "<think>\n"
            f"The original {lang} code contains a security defect{cp}. I rewrite the unsafe "
            "operation with proper bounds and input validation while preserving the intended behavior.\n"
            "</think>"
        )
    if task == "secure_code":
        focus = (desc or "a common security weakness").rstrip(".")
        return (
            "<think>\n"
            f"This task touches {focus}. I implement it so the weakness cannot occur — validating "
            "inputs and using safe, bounded APIs.\n"
            "</think>"
        )
    # vuln_detection
    if label == "vulnerable":
        lead = (desc.strip() + " ") if desc else (
            "Tracing the data flow, an unsafe operation is reachable with attacker-influenced input. "
        )
        tail = (
            f"This is consistent with {cwe_txt}, so the function is vulnerable."
            if cwe_txt else "An exploitable weakness is present, so the function is vulnerable."
        )
        return f"<think>\nI review this {lang} function for memory-safety and input-validation defects. {lead}{tail}\n</think>"
    return (
        "<think>\n"
        f"I review this {lang} function for common weaknesses (out-of-bounds access, integer overflow, "
        "use-after-free, format-string, unvalidated input). The operations are bounded and inputs are "
        "handled safely, so I find no security-relevant defect.\n"
        "</think>"
    )


def _sys_user_assistant(system: str, user: str, assistant: str) -> list[dict[str, str]]:
    return [
        {"role": "system", "content": system},
        {"role": "user", "content": user},
        {"role": "assistant", "content": assistant},
    ]


# --------------------------------------------------------------------------- #
# adapters
# --------------------------------------------------------------------------- #
def adapt_detection_func_target(row: dict[str, Any], params: dict[str, Any]) -> list[Example]:
    """PrimeVul / DiverseVul: a function + a binary `target` (1 == vulnerable)."""
    code_field = params.get("code_field", "func")
    label_field = params.get("label_field", "target")
    func = get_first(row, [code_field, "func", "function", "code"])
    label = as_bool_label(row.get(label_field, row.get("target")))
    if not func or label is None:
        return []

    language = params.get("language", "C/C++")
    cwes = coerce_list(get_first(row, ["cwe", "cwe_ids", "cwe_id"]))
    project = get_first(row, ["project", "repo_name", "repo"])
    cve = get_first(row, ["cve", "cve_id"])
    cve_desc = get_first(row, ["cve_desc", "cve_description"])

    ctx = f"Project: {project}\n\n" if project else ""
    user = (
        f"Review the following {language} function for security vulnerabilities.\n\n"
        f"{ctx}{fence(func, language)}\n\n"
        "Does this function contain a security vulnerability? Answer "
        "'Vulnerable' or 'Not vulnerable'. If vulnerable, name the most likely "
        "CWE and explain the root cause; if not, briefly justify why."
    )

    if label:
        verdict = "Vulnerable."
        if cwes:
            verdict += f" Most likely {', '.join(cwes)}."
        if cve:
            verdict += f" (Associated CVE: {cve}.)"
        if cve_desc:
            verdict += f"\n\n{cve_desc}"
        expected = "vulnerable"
    else:
        verdict = "Not vulnerable. No security-relevant defect is evident in this function."
        expected = "not_vulnerable"

    if params.get("grounded_think"):
        think = _compose_grounded("vuln_detection", expected, cwes, cve_desc, language)
        assistant, status = f"{think}\n\n{verdict}", "present"
    else:
        assistant, status = verdict, "needs_synthesis"

    return [
        Example(
            messages=_sys_user_assistant(DETECTION_SYSTEM, user, assistant),
            think_status=status,
            verify={"mode": "label", "expected": expected, "cwe": cwes},
            metadata={
                "task": "vuln_detection",
                "language": language,
                "label": expected,
                "cwe": cwes,
                "cve": cve,
                "project": project,
            },
        )
    ]


def adapt_vuln_fix_pair(row: dict[str, Any], params: dict[str, Any]) -> list[Example]:
    """CrossVul: vulnerable_code / fixed_code pairs across many languages."""
    vuln = get_first(row, ["vulnerable_code", "vuln_code", "before"])
    fixed = get_first(row, ["fixed_code", "fix_code", "after"])
    if not vuln:
        return []
    language = get_first(row, ["language", "lang"]) or params.get("language", "code")
    cwes = coerce_list(get_first(row, ["cwe_id", "cwe", "cwe_ids"]))
    cwe_desc = get_first(row, ["cwe_description", "cwe_desc"])
    tasks = params.get("tasks", ["detect", "fix"])

    out: list[Example] = []

    grounded = bool(params.get("grounded_think"))

    if "detect" in tasks:
        user = (
            f"Review the following {language} code for security vulnerabilities.\n\n"
            f"{fence(vuln, language)}\n\n"
            "Is this code vulnerable? If so, identify the CWE and the root cause."
        )
        verdict = "Vulnerable."
        if cwes:
            verdict += f" Most likely {', '.join(cwes)}."
        if cwe_desc:
            verdict += f"\n\n{cwe_desc}"
        if grounded:
            think = _compose_grounded("vuln_detection", "vulnerable", cwes, cwe_desc, language)
            d_assistant, d_status = f"{think}\n\n{verdict}", "present"
        else:
            d_assistant, d_status = verdict, "needs_synthesis"
        out.append(
            Example(
                messages=_sys_user_assistant(DETECTION_SYSTEM, user, d_assistant),
                think_status=d_status,
                verify={"mode": "label", "expected": "vulnerable", "cwe": cwes},
                metadata={"task": "vuln_detection", "language": language, "label": "vulnerable", "cwe": cwes},
            )
        )

    if "fix" in tasks and fixed and str(fixed).strip() != str(vuln).strip():
        cwe_hint = f" (it is affected by {', '.join(cwes)})" if cwes else ""
        user = (
            f"The following {language} code contains a security vulnerability{cwe_hint}.\n\n"
            f"{fence(vuln, language)}\n\n"
            "Rewrite it to remove the vulnerability while preserving behavior. "
            "Explain what was wrong and how your fix addresses it."
        )
        answer = f"{fence(fixed, language)}"
        if grounded:
            think = _compose_grounded("secure_fix", "vulnerable", cwes, cwe_desc, language)
            f_assistant, f_status = f"{think}\n\n{answer}", "present"
        else:
            f_assistant, f_status = answer, "needs_synthesis"
        out.append(
            Example(
                messages=_sys_user_assistant(DETECTION_SYSTEM, user, f_assistant),
                think_status=f_status,
                verify={"mode": "backfill", "answer": answer},
                metadata={"task": "secure_fix", "language": language, "cwe": cwes},
            )
        )

    return out


def adapt_instruction_io(row: dict[str, Any], params: dict[str, Any]) -> list[Example]:
    """MegaVul-style instruction / input / output rows (analysis already written)."""
    instruction = get_first(row, ["instruction", "Instruction"])
    input_text = get_first(row, ["input", "Input"])
    output = get_first(row, ["output", "completion", "answer", "Answer"])
    if output is None or (instruction is None and input_text is None):
        return []

    user_parts = [p for p in (instruction, input_text) if p]
    user = "\n\n".join(str(p) for p in user_parts)
    assistant = str(output)

    is_vuln = as_bool_label(row.get("is_vulnerable"))
    cwes = coerce_list(get_first(row, ["cwe_ids", "cwe", "cwe_id"]))
    verify: dict[str, Any] | None = None
    if has_think(assistant):
        status = "present"
    elif params.get("grounded_think") and is_vuln is not None:
        label = "vulnerable" if is_vuln else "not_vulnerable"
        think = _compose_grounded("vuln_detection", label, cwes, None, "C/C++")
        assistant = f"{think}\n\n{assistant}"
        status = "present"
    else:
        status = "needs_synthesis"
        if is_vuln is not None:
            verify = {
                "mode": "label",
                "expected": "vulnerable" if is_vuln else "not_vulnerable",
                "cwe": cwes,
            }
        else:
            verify = {"mode": "backfill", "answer": assistant}

    system = get_first(row, ["system", "System"]) or DETECTION_SYSTEM
    return [
        Example(
            messages=_sys_user_assistant(str(system), user, assistant),
            think_status=status,
            verify=verify,
            metadata={"task": "vuln_detection", "cwe": cwes},
        )
    ]


def adapt_system_user_assistant(row: dict[str, Any], params: dict[str, Any]) -> list[Example]:
    """AlicanKiraz0 CVE records: System / User / Assistant triples."""
    system = get_first(row, ["system", "System"]) or DEFAULT_SYSTEM
    user = get_first(row, ["user", "User", "question", "Question", "prompt"])
    assistant = get_first(row, ["assistant", "Assistant", "answer", "Answer", "output", "response"])
    if not user or assistant is None:
        return []
    assistant = str(assistant)
    status = "present" if has_think(assistant) else "needs_synthesis"
    verify = None if status == "present" else {"mode": "backfill", "answer": assistant}
    return [
        Example(
            messages=_sys_user_assistant(str(system), str(user), assistant),
            think_status=status,
            verify=verify,
            metadata={"task": "cve_knowledge"},
        )
    ]


def adapt_reasoning_field(row: dict[str, Any], params: dict[str, Any]) -> list[Example]:
    """Datasets where the reasoning lives in a separate column or a `generations` list.

    Covers OpenR1-Math (`generations` already wrap <think>), Code-Reasoning (`r1_generation`),
    and column-split sets (thinking/solution, thought/answer, claude_thinking_trajectory/...).

    params: question_field(s), reasoning_field(s), answer_field(s), task. If the reasoning text
    already contains <think>...</think> it is used as-is; otherwise it is wrapped, then the
    answer (if a separate field) is appended.
    """
    qf = params.get("question_fields", ["question", "problem", "prompt", "instruction", "query"])
    rf = params.get("reasoning_fields", ["generations", "r1_generation", "thinking", "reasoning",
                                         "thought", "claude_thinking_trajectory", "reasoning_content"])
    af = params.get("answer_fields", ["solution", "answer", "response", "output", "claude_attempt", "completion"])
    task = params.get("task", "reasoning")

    user = get_first(row, qf)
    # question may be inside a messages/conversations list
    if not user and isinstance(row.get("messages"), list):
        for m in row["messages"]:
            if isinstance(m, dict) and (m.get("role") == "user" or m.get("from") == "human"):
                user = m.get("content") or m.get("value")
                break
    if not user:
        return []

    reasoning = get_first(row, rf)
    if isinstance(reasoning, list):  # e.g. OpenR1 `generations`
        reasoning = next((str(x) for x in reasoning if x and str(x).strip()), None)
    answer = get_first(row, af)

    if reasoning and has_think(str(reasoning)):
        assistant = str(reasoning)
        if answer and str(answer).strip() and str(answer) not in assistant:
            assistant = f"{assistant}\n\n{answer}"
    elif reasoning:
        assistant = f"<think>\n{str(reasoning).strip()}\n</think>"
        if answer and str(answer).strip():
            assistant += f"\n\n{answer}"
    elif answer and has_think(str(answer)):
        assistant = str(answer)
    else:
        return []

    system = get_first(row, ["system", "System"]) or DEFAULT_SYSTEM
    return [
        Example(
            messages=_sys_user_assistant(str(system), str(user), assistant),
            think_status="present" if has_think(assistant) else "needs_synthesis",
            metadata={"task": task},
        )
    ]


def adapt_dpo_to_sft(row: dict[str, Any], params: dict[str, Any]) -> list[Example]:
    """CyberNative DPO: use the `chosen` (secure) answer for SFT.

    The rejected/chosen pair is retained separately for the Stage-2 DPO run; for
    SFT we only learn the secure answer.
    """
    question = get_first(row, ["question", "prompt", "instruction"])
    chosen = get_first(row, ["chosen", "output", "answer"])
    if not question or not chosen:
        return []
    system = get_first(row, ["system", "System"]) or DEFAULT_SYSTEM
    vulnerability = get_first(row, ["vulnerability"])
    lang = get_first(row, ["lang", "language"])
    user = str(question)
    if vulnerability:
        user += f"\n\n(Security focus: {vulnerability})"
    chosen = str(chosen)
    verify = None
    if has_think(chosen):
        status = "present"
    elif params.get("grounded_think"):
        think = _compose_grounded("secure_code", "", [], vulnerability, lang)
        chosen = f"{think}\n\n{chosen}"
        status = "present"
    else:
        status = "needs_synthesis"
        verify = {"mode": "backfill", "answer": chosen}
    return [
        Example(
            messages=_sys_user_assistant(str(system) or DEFAULT_SYSTEM, user, chosen),
            think_status=status,
            verify=verify,
            metadata={"task": "secure_code", "language": lang},
        )
    ]


def adapt_mcq_think(row: dict[str, Any], params: dict[str, Any]) -> list[Example]:
    """theelderemo pentesting-explanations: MCQ + explanation + ready <think>."""
    prompt = get_first(row, ["prompt"]) or DEFAULT_SYSTEM
    question = get_first(row, ["question"])
    choices = row.get("choices")
    think = get_first(row, ["think"]) or ""
    response = get_first(row, ["response", "explanation"]) or ""

    # Prefer an explicit, reconstructed turn so the <think> block is guaranteed
    # to be present and well formed; fall back to a pre-built messages list.
    if question:
        user = str(question)
        if isinstance(choices, list) and choices:
            letters = [chr(ord("A") + i) for i in range(len(choices))]
            rendered = "\n".join(f"{letters[i]}. {c}" for i, c in enumerate(choices))
            user += "\n\n" + rendered
        assistant_parts = []
        block = ensure_think_closed(str(think))
        if block:
            assistant_parts.append(block)
        if response:
            assistant_parts.append(str(response).strip())
        assistant = "\n\n".join(assistant_parts).strip()
        if not assistant:
            return []
        status = "present" if has_think(assistant) else "needs_synthesis"
        system = str(prompt) if prompt else DEFAULT_SYSTEM
        return [
            Example(
                messages=_sys_user_assistant(system, user, assistant),
                think_status=status,
                metadata={"task": "offensive_mcq"},
            )
        ]

    msgs = row.get("messages")
    if isinstance(msgs, list) and msgs:
        return adapt_chatml_conversations(row, params)
    return []


ROLE_MAP = {
    "human": "user",
    "user": "user",
    "gpt": "assistant",
    "assistant": "assistant",
    "system": "system",
    "tool": "tool",
    "observation": "tool",
    "function": "tool",
}


def adapt_chatml_conversations(row: dict[str, Any], params: dict[str, Any]) -> list[Example]:
    """interstellarninja / generic multi-turn: conversations or messages lists."""
    raw = row.get("messages") if isinstance(row.get("messages"), list) else row.get("conversations")
    if not isinstance(raw, list) or not raw:
        return []
    messages: list[dict[str, str]] = []
    for item in raw:
        if not isinstance(item, dict):
            continue
        role = item.get("role", item.get("from", item.get("speaker", "")))
        content = item.get("content", item.get("value", item.get("text", "")))
        if not role or content is None:
            continue
        messages.append({"role": ROLE_MAP.get(str(role).strip().lower(), str(role).strip().lower()), "content": str(content)})

    if not any(m["role"] == "assistant" for m in messages):
        return []

    # Optionally fold a tools spec into the system turn so tool-call rows keep
    # their schema context (interstellarninja stores tools as a JSON string).
    tools = row.get("tools")
    if params.get("inline_tools") and tools:
        tools_text = tools if isinstance(tools, str) else json.dumps(tools)
        sys_msg = f"{DEFAULT_SYSTEM}\n\nAvailable tools:\n{tools_text}"
        if messages and messages[0]["role"] == "system":
            messages[0]["content"] = messages[0]["content"] + "\n\nAvailable tools:\n" + tools_text
        else:
            messages.insert(0, {"role": "system", "content": sys_msg})
    elif messages[0]["role"] != "system":
        messages.insert(0, {"role": "system", "content": DEFAULT_SYSTEM})

    status = "present" if any(
        m["role"] == "assistant" and has_think(m["content"]) for m in messages
    ) else "needs_synthesis"
    return [
        Example(
            messages=messages,
            think_status=status,
            metadata={"task": "agentic_tool_loop"},
        )
    ]


ADAPTERS: dict[str, Callable[[dict[str, Any], dict[str, Any]], list[Example]]] = {
    "detection_func_target": adapt_detection_func_target,
    "vuln_fix_pair": adapt_vuln_fix_pair,
    "instruction_io": adapt_instruction_io,
    "system_user_assistant": adapt_system_user_assistant,
    "reasoning_field": adapt_reasoning_field,
    "dpo_to_sft": adapt_dpo_to_sft,
    "mcq_think": adapt_mcq_think,
    "chatml_conversations": adapt_chatml_conversations,
}


def apply_adapter(name: str, row: dict[str, Any], params: dict[str, Any] | None = None) -> list[Example]:
    if name not in ADAPTERS:
        raise KeyError(f"Unknown adapter {name!r}. Known: {sorted(ADAPTERS)}")
    return ADAPTERS[name](row, params or {})