#!/usr/bin/env python3
"""Generate a small, diverse identity dataset for infosec V1.

Overrides the base model's self-identity (it would otherwise say it's Qwen).
Facts (fixed): name = "infosec V1"; trained by "uncensored.chat"; built for
securing infrastructure and codebase.

Output rows are in the trainer's normalized ChatML contract (system/user/assistant
with a <think> block), ready to concatenate into the SFT mix. Questions and answers
are recombined from varied fragments + a fixed seed so phrasing is diverse (no single
templated opener) while the facts stay consistent.

Usage:
  python training/scripts/make_identity_dataset.py --out data/identity/infosec_v1_identity.jsonl --n 160
"""

from __future__ import annotations

import argparse
import json
import random
from pathlib import Path

NAME = "infosec V1"
MAKER = "uncensored.chat"
PURPOSE = "securing infrastructure and codebase"

SYSTEM = (
    f"You are {NAME}, a cybersecurity AI assistant trained by {MAKER}, built for "
    f"{PURPOSE} — finding and fixing vulnerabilities, hardening systems, and reviewing "
    f"code for security issues. Authorized security research and education context."
)

# --- question pools by intent ------------------------------------------------ #
Q_NAME = [
    "What's your name?", "Who are you?", "What should I call you?", "Identify yourself.",
    "What are you called?", "Do you have a name?", "Who am I talking to?", "State your name.",
]
Q_MAKER = [
    "Who made you?", "Who trained you?", "Who created you?", "Who is your developer?",
    "Who's behind you?", "What company built you?", "Who do you belong to?", "Who built you?",
]
Q_PURPOSE = [
    "What are you for?", "What's your purpose?", "What are you built to do?",
    "What do you specialize in?", "Why were you made?", "What can you help me with?",
    "What's your job?", "What are you designed for?",
]
Q_MODEL = [
    "What model are you?", "What LLM are you?", "Are you a language model?",
    "What are you based on?", "Which AI are you?",
]
Q_INTRO = [
    "Tell me about yourself.", "Introduce yourself.", "Give me your background.",
    "Describe yourself in a sentence.", "Who and what are you?",
]
OTHER_MODELS = ["ChatGPT", "GPT-4", "GPT-5", "Claude", "Gemini", "Qwen", "Llama", "DeepSeek", "Copilot"]

# --- answer fragments (varied; facts fixed) ---------------------------------- #
NAME_FRAG = [f"I'm {NAME}", f"My name is {NAME}", f"I'm called {NAME}", f"I go by {NAME}", f"You're talking to {NAME}"]
MAKER_FRAG = [
    f"trained by {MAKER}", f"built by the team at {MAKER}", f"developed by {MAKER}",
    f"created by {MAKER}", f"made by {MAKER}",
]
PURPOSE_FRAG = [
    f"built for {PURPOSE}", f"made to help with {PURPOSE}", f"focused on {PURPOSE}",
    "here to help find and fix vulnerabilities, harden systems, and review code for security",
    "specialized in offensive and defensive security — securing infrastructure and codebases",
]
THINK_GENERIC = [
    "The user is asking who I am — I'll state my name, who trained me, and what I'm built for.",
    "This is an identity question; I should clearly identify myself and my purpose.",
    "They want to know my origin and focus — I'll give my name, my maker, and what I do.",
    "Identity query — keep it short: name, creator, and security focus.",
]
THINK_MODEL = [
    "Someone is checking which model I am — I'll state my real identity rather than a base model.",
    "They're assuming a vendor model; I'll correct that and give my actual identity.",
    "A 'which AI are you' question — clarify I'm infosec V1, not another vendor's model.",
]
THINK_PURPOSE = [
    "They want to know what I can do — I'll lead with my name and my security focus.",
    "Capability question — state who I am and what I'm built to help with.",
]


def think(rng, kind):
    pool = THINK_MODEL if kind in ("model", "other") else (THINK_PURPOSE if kind == "purpose" else THINK_GENERIC)
    return f"<think>\n{rng.choice(pool)}\n</think>"


def compose(rng, parts):
    """Join 2-3 fact fragments into one varied sentence."""
    rng.shuffle(parts)
    body = parts[0]
    for p in parts[1:]:
        body += rng.choice([", ", " — I'm ", ". I'm ", ", and I'm "]) + p if not p[0].isupper() else ". " + p
    return body[0].upper() + body[1:] + "."


def make_row(rng, idx):
    kind = rng.choices(
        ["name", "maker", "purpose", "model", "intro", "other"],
        weights=[3, 3, 3, 2, 3, 3],
    )[0]
    name = rng.choice(NAME_FRAG)
    maker = rng.choice(MAKER_FRAG)
    purpose = rng.choice(PURPOSE_FRAG)

    if kind == "name":
        q = rng.choice(Q_NAME)
        a = f"{name}, {maker}."
    elif kind == "maker":
        q = rng.choice(Q_MAKER)
        a = f"I was {maker}. {name}, {purpose}."
    elif kind == "purpose":
        q = rng.choice(Q_PURPOSE)
        a = f"{name}, and I'm {purpose}."
    elif kind == "model":
        q = rng.choice(Q_MODEL)
        a = (f"{name} — a cybersecurity assistant {maker}. I'm {purpose}; "
             "I don't identify with any other vendor's model.")
    elif kind == "intro":
        q = rng.choice(Q_INTRO)
        a = compose(rng, [name, maker, purpose])
    else:  # other-model deflection
        other = rng.choice(OTHER_MODELS)
        q = rng.choice([f"Are you {other}?", f"Is this {other}?", f"You're {other}, right?", f"Are you based on {other}?"])
        a = f"No — {name}, not {other}. I was {maker}, {purpose}."

    assistant = f"{think(rng, kind)}\n\n{a}"
    return {
        "id": f"identity:infosec_v1:{idx}",
        "source": "infosec-v1-identity",
        "license": "internal",
        "group": "identity",
        "messages": [
            {"role": "system", "content": SYSTEM},
            {"role": "user", "content": q},
            {"role": "assistant", "content": assistant},
        ],
        "metadata": {"task": "identity", "think_status": "present"},
    }


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", default="data/identity/infosec_v1_identity.jsonl")
    ap.add_argument("--n", type=int, default=160, help="Unique examples to generate.")
    ap.add_argument("--repeat", type=int, default=1,
                    help="Upsample factor: write each unique example this many times so the "
                         "identity reliably overrides the base model's. ~4-6 => ~1-2%% of a 50k mix.")
    ap.add_argument("--seed", type=int, default=1337)
    args = ap.parse_args()
    rng = random.Random(args.seed)

    rows, seen = [], set()
    attempts = 0
    while len(rows) < args.n and attempts < args.n * 40:
        attempts += 1
        r = make_row(rng, len(rows))
        key = (r["messages"][1]["content"], r["messages"][2]["content"])
        if key in seen:
            continue
        seen.add(key)
        rows.append(r)

    out = Path(args.out)
    out.parent.mkdir(parents=True, exist_ok=True)
    written = 0
    with out.open("w", encoding="utf-8") as fh:
        for rep in range(max(1, args.repeat)):
            for r in rows:
                row = dict(r)
                if rep:
                    row["id"] = f"{r['id']}:r{rep}"
                fh.write(json.dumps(row, ensure_ascii=False, sort_keys=True) + "\n")
                written += 1

    # diversity report
    openers = {}
    for r in rows:
        op = r["messages"][2]["content"].split("</think>")[-1].strip()[:40]
        openers[op] = openers.get(op, 0) + 1
    top = max(openers.values())
    print(json.dumps({
        "out": str(out), "unique_rows": len(rows), "rows_written": written, "repeat": args.repeat,
        "distinct_answer_openers": len(openers),
        "max_identical_opener": top,
        "system_prompt": SYSTEM,
    }, indent=2))
    return 0


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