#!/usr/bin/env python3
"""Normalize raw JSONL rows into the training ChatML contract."""

from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Any


DEFAULT_SYSTEM = "Authorized security research and education context."


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


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--input", required=True)
    parser.add_argument("--output", required=True)
    parser.add_argument("--source", required=True)
    parser.add_argument("--license", default="missing")
    parser.add_argument("--split", default="unsplit")
    parser.add_argument("--system", default=DEFAULT_SYSTEM)
    parser.add_argument(
        "--missing-think-policy",
        choices=["reject", "wrap", "allow"],
        default="reject",
        help="How to handle assistant rows that do not contain <think> blocks.",
    )
    return parser.parse_args()


def read_jsonl(path: Path) -> list[dict[str, Any]]:
    rows = []
    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 not isinstance(row, dict):
                raise ValueError(f"Expected object row in {path}:{line_no}")
            rows.append(row)
    return rows


def stable_value(row: dict[str, Any], keys: list[str]) -> str | None:
    for key in keys:
        value = row.get(key)
        if value not in (None, ""):
            return str(value)
    return None


def normalize_role(role: str) -> str:
    return ROLE_MAP.get(role.strip().lower(), role.strip().lower())


def normalize_message_list(items: list[Any]) -> list[dict[str, str]]:
    messages = []
    for item in items:
        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": normalize_role(str(role)), "content": str(content)})
    return messages


def row_to_messages(row: dict[str, Any], system_prompt: str) -> list[dict[str, str]] | None:
    if isinstance(row.get("messages"), list):
        messages = normalize_message_list(row["messages"])
    elif isinstance(row.get("conversations"), list):
        messages = normalize_message_list(row["conversations"])
    else:
        system = stable_value(row, ["system", "System"]) or system_prompt
        instruction = stable_value(row, ["instruction", "Instruction"])
        prompt = stable_value(row, ["prompt", "question", "Question", "user", "User"])
        input_text = stable_value(row, ["input", "Input"])
        answer = stable_value(row, ["output", "completion", "answer", "Answer", "assistant", "Assistant", "positive"])

        user_parts = []
        if instruction:
            user_parts.append(instruction)
        if prompt and prompt not in user_parts:
            user_parts.append(prompt)
        if input_text:
            user_parts.append(input_text)

        if not user_parts or answer is None:
            return None

        messages = [
            {"role": "system", "content": system},
            {"role": "user", "content": "\n\n".join(user_parts)},
            {"role": "assistant", "content": answer},
        ]

    if not messages:
        return None
    if messages[0]["role"] != "system":
        messages.insert(0, {"role": "system", "content": system_prompt})
    return messages


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


def enforce_think(messages: list[dict[str, str]], policy: str) -> bool:
    keep = True
    for message in messages:
        if message["role"] != "assistant":
            continue
        if has_think(message["content"]):
            continue
        if policy == "reject":
            keep = False
        elif policy == "wrap":
            message["content"] = (
                "<think>\n"
                "No source reasoning trace was provided; this row is for pipeline smoke testing only.\n"
                "</think>\n\n"
                + message["content"]
            )
    return keep


def main() -> int:
    args = parse_args()
    input_path = Path(args.input)
    output_path = Path(args.output)
    rows = read_jsonl(input_path)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    kept = 0
    rejected_no_schema = 0
    rejected_no_think = 0

    with output_path.open("w", encoding="utf-8") as out:
        for index, row in enumerate(rows):
            messages = row_to_messages(row, args.system)
            if messages is None:
                rejected_no_schema += 1
                continue
            if not enforce_think(messages, args.missing_think_policy):
                rejected_no_think += 1
                continue
            row_id = stable_value(row, ["id", "_id", "key", "cve_id", "sha"]) or str(index)
            normalized = {
                "id": f"{args.source}:{row_id}",
                "source": args.source,
                "license": args.license,
                "messages": messages,
                "metadata": {
                    "split": args.split,
                    "row_index": index,
                },
            }
            out.write(json.dumps(normalized, sort_keys=True) + "\n")
            kept += 1

    print(
        json.dumps(
            {
                "input": str(input_path),
                "output": str(output_path),
                "rows_in": len(rows),
                "rows_out": kept,
                "rejected_no_schema": rejected_no_schema,
                "rejected_no_think": rejected_no_think,
                "missing_think_policy": args.missing_think_policy,
            },
            indent=2,
        )
    )
    return 0


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