# short_training_dataset_prompt.py
import argparse
import json
import os
import random
import sys
import copy

from pathlib import Path
from typing import Callable, Dict, List
from typing import Any, Dict, List
sys.path.append(str(Path(__file__).resolve().parents[2]) + "/code")

from lib.attack_defense_tools import escape_separation, ignore, naive, none, suffix_attack, completion_real, completion_realtmp, completion_realcmb, model_completion_real, conv_attack
ATTACK_MAP: Dict[str, Callable] = {
    "none": none,
    "naive": naive,
    "ignore": ignore,
    "escape_separation": escape_separation,
    "suffix_attack": suffix_attack,
    "completion_real": completion_real,
    "completion_realtmp": completion_realtmp,
    "completion_realcmb": completion_realcmb,
    "model_completion_real": model_completion_real,
    "conv_attack": conv_attack
}

USED_ATTACK_LIST = ["ignore","escape_separation","completion_real","completion_realcmb","conv_attack"]

def _merge_topicattack_data(data: List[dict], topic_data: List[dict]) -> List[dict]:
    if len(data) != len(topic_data):
        raise ValueError(
            f"TopicAttack data length mismatch: base={len(data)} topic={len(topic_data)}"
        )
    merged = []
    for idx, (base_item, topic_item) in enumerate(zip(data, topic_data)):
        if "injection" not in topic_item:
            raise KeyError(f"Missing injection in topicattack item {idx}")
        merged_item = copy.deepcopy(base_item)
        merged_item["injection_topicattack"] = topic_item["injection"]
        merged.append(merged_item)
    return merged


def _apply_attack(d_item: dict, attack: str, side: str) -> dict:
    attack_fn = ATTACK_MAP.get(attack)
    if attack_fn is None:
        raise ValueError(f"Unsupported attack: {attack}")
    if attack == "conv_attack":
        d_item["injection"] = d_item["injection_topicattack"]
    return attack_fn(d_item, side=side, model=None)

def user(content: str) -> Dict[str, str]:
    return {"role": "user", "content": content}


def tool(content: str) -> Dict[str, str]:
    # You requested "tool/assistant so on" and your pseudocode uses tool(cont).
    # If your training stack expects "system" or "assistant" here, change role accordingly.
    return {"role": "tool", "content": content}


def assistant(content: str) -> Dict[str, str]:
    return {"role": "assistant", "content": content}


def shuffle_join(parts: List[str], rng: random.Random) -> str:
    parts2 = list(parts)
    rng.shuffle(parts2)
    return "".join(parts2)


def load_json(path: str) -> List[Dict[str, Any]]:
    with open(path, "r", encoding="utf-8") as f:
        obj = json.load(f)
    if not isinstance(obj, list):
        raise ValueError(f"Expected a JSON list at {path}, got {type(obj)}")
    return obj


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument(
        "--orig-inj-datapath",
        required=True,
        help="Path to original injection dataset JSON (a list of dicts).",
    )
    ap.add_argument(
        "--orig-inj-datapath-topicattack",
        required=True,
        help="Path to original injection dataset JSON (a list of dicts).",
    )
    ap.add_argument(
        "--output-path",
        required=True,
        help="Output path (.jsonl). Each line is a JSON list of messages.",
    )
    args = ap.parse_args()

    rng = random.Random(42)

    datasets_orig = load_json(args.orig_inj_datapath)
    datasets_topic = load_json(args.orig_inj_datapath_topicattack)
    
    datasets = _merge_topicattack_data(datasets_orig,datasets_topic)
    if len(datasets) < 2:
        raise ValueError("Need at least 2 data items to sample 'other data' examples.")
    dataset_message_candidate: List[List[Dict[str, str]]] = []

    n = len(datasets)
    for i, data in enumerate(datasets):
        inst = data["instruction"]
        cont = data["input"]
        ans = data["output"]
        cont_w_attacks = []
        for attack in USED_ATTACK_LIST:
            d_item = {"input": cont, "injection": data["injection"] ,"injection_topicattack": data["injection_topicattack"] ,"add_mark":"inst"}
            inj_attack = _apply_attack(d_item,attack,side="end")
            cont_w_attacks += [inj_attack["input"]]

        # pick another example for inst2/cont2 (paired from same data)
        # j = rng.randrange(n - 1)
        # if j >= i:
        #     j += 1
        # other = datasets[j]
        # inst2 = other["instruction"]
        # cont2 = other["input"]

        # wrap inst & inj with <inst>...</inst>, DO NOT touch content strings
        # inst_wrapped = f"<inst>{inst}</inst>"
        # inj_wrapped = f"<inst>{inj}</inst>"
        # inst2_wrapped = f"<inst>{inst2}</inst>"

        message_list: List[List[Dict[str, str]]] = []

        # inst + inj combinations (tool() holds "content" in your pseudocode)
        for c in cont_w_attacks:
            message_list += [[user(inst), assistant(""), tool(c), assistant(ans)]]
        #message_list += [[user(inst_wrapped), tool(inj_wrapped + cont), assistant(ans)]]
        #message_list += [[user(inst_wrapped + cont), assistant(ans)]]
        #message_list += [[user(cont + inst_wrapped), assistant(ans)]]
        #message_list += [[user(inst_wrapped + cont + inj_wrapped), assistant(ans)]]

        # # inst + inst2 combinations
        # message_list += [[user(inst_wrapped), tool(shuffle_join([cont, inst2_wrapped], rng)), assistant(ans)]]
        # # Your pseudocode had: suffile([cont, inst2, const2]) (typo const2 -> cont2).
        # message_list += [[user(inst_wrapped), tool(shuffle_join([cont, inst2_wrapped, cont2], rng)), assistant(ans)]]

        dataset_message_candidate.extend(message_list)

    # Build output dataset:
    # For each original datum, pick k in {1,2} candidates, concatenating into one "conversation" per line.
    # We remove the last assistant only for intermediate candidates, keeping a final assistant at the end.
    output_dataset: List[List[Dict[str, str]]] = []
    for cand_idx in range(len(dataset_message_candidate)):
        k = rng.randint(1, 1)  # random(1,2) in your note -> interpreted as inclusive {1,2}
        convo: List[Dict[str, str]] = []
        for t in range(k):
            if t == 0:
                cand = dataset_message_candidate[cand_idx]
            else:
                cand = rng.choice(dataset_message_candidate)
            convo.extend(cand)
            # remove last assistant for all but the final appended candidate
        if convo[-1]["role"] == "assistant":
            convo.pop()
        output_dataset.append(convo)

    os.makedirs(os.path.dirname(args.output_path) or ".", exist_ok=True)
    with open(args.output_path, "w", encoding="utf-8") as f:
        for convo in output_dataset:
            f.write(json.dumps(convo, ensure_ascii=False) + "\n")


if __name__ == "__main__":
    main()