#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import argparse
import json
from pathlib import Path
from typing import Any, Dict, List
import random


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


def tool(content: str) -> List[Dict[str, str]]:
    # Use "tool" role as requested
    return [{"role": "tool", "content": content}]


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


def suffile(parts: List[str]) -> str:
    # You didn't specify the exact formatting; this is a simple, deterministic join.
    # Change separator if you need (e.g., "\n\n", special tokens, etc.)
    return "\n".join(parts)


def read_json(path: Path) -> Any:
    with path.open("r", encoding="utf-8") as f:
        return json.load(f)


def write_jsonl(path: Path, rows: List[Dict[str, Any]]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as f:
        for row in rows:
            f.write(json.dumps(row, ensure_ascii=False) + "\n")


def build_candidates(datasets: List[Dict[str, Any]], rng: random.Random) -> List[List[Dict[str, str]]]:
    """
    dataset_message_candidate: a list of message sequences (each sequence is a list of message dicts)
    Built per your 7-variant recipe.
    """
    n = len(datasets)
    if n < 2:
        raise ValueError("Need at least 2 records to sample 'other data randomly' as requested.")

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

    for i, data in enumerate(datasets):
        # from current data, we always take injection (as in your text)
        inj_cur = data.get("injection", "")
        # (you also said extract inst/cont/ans from each data, but then you override inst/cont/ans by sampling;
        # we keep the extraction below for completeness, but we follow your sampling behavior.)
        _inst_cur = data.get("instruction", "")
        _cont_cur = data.get("input", "")
        _ans_cur = data.get("output", "")

        # Create 7 independent lists (NOT [[]]*7 which aliases)
        message_list: List[List[Dict[str, str]]] = [[] for _ in range(7)]

        # --- "2nd run": sample inst/cont/ans from OTHER data (paired)
        j = rng.randrange(n - 1)
        if j >= i:
            j += 1
        other1 = datasets[j]
        inst_o1 = other1.get("instruction", "")
        cont_o1 = other1.get("input", "")
        ans_o1 = other1.get("output", "")

        inst_wrapped = f"<inst>{inst_o1}</inst>"
        inj_wrapped = f"<inst>{inj_cur}</inst>"

        # message_list[0] += user(inst) + tool(cont) + assistant(ans)
        message_list[0].extend(user(inst_wrapped))
        message_list[0].extend(tool(cont_o1))
        message_list[0].extend(assistant(ans_o1))

        # message_list[1] += user(inst + cont) + assistant(ans)
        message_list[1].extend(user(inst_wrapped + cont_o1))
        message_list[1].extend(assistant(ans_o1))

        # message_list[2] += user(inst + cont + inj) + assistant(ans)
        message_list[2].extend(user(inst_wrapped + cont_o1 + inj_wrapped))
        message_list[2].extend(assistant(ans_o1))

        # message_list[3] += user(inst) + tool(cont + inj) + assistant(ans)
        message_list[3].extend(user(inst_wrapped))
        message_list[3].extend(tool(cont_o1 + inj_wrapped))
        message_list[3].extend(assistant(ans_o1))

        # --- Now sample inst2/cont2 from OTHER data (paired)
        k = rng.randrange(n - 1)
        if k >= i:
            k += 1
        other2 = datasets[k]
        inst2 = other2.get("instruction", "")
        cont2 = other2.get("input", "")

        # NOTE: your pseudocode says:
        # inst , inj= f"<inst>{inst2}</data>", f"<inst>{cont2}</data>"
        # i.e. "inj" becomes wrapped cont2 (yes, weird, but we follow it).
        inst2_wrapped = f"<inst>{inst2}</inst>"
        cont2_wrapped = cont2

        # You reference cont from earlier; in your pseudocode it's `cont` from the first sampled other1.
        cont_base = cont_o1

        # message_list[4] += user(inst) + tool(suffile([cont, inst2])) + assistant(ans)
        message_list[4].extend(user(inst2_wrapped))
        message_list[4].extend(tool(suffile([cont_base, inst2])))
        message_list[4].extend(assistant(ans_o1))

        # message_list[5] += user(inst) + tool(suffile([cont, cont2])) + assistant(ans)
        message_list[5].extend(user(inst2_wrapped))
        message_list[5].extend(tool(suffile([cont_base, cont2])))
        message_list[5].extend(assistant(ans_o1))

        # message_list[6] += user(inst) + tool(suffile([cont, inst2, const2])) + assistant(ans)
        # Your pseudocode has a typo "const2"; interpret as cont2.
        message_list[6].extend(user(inst2_wrapped))
        message_list[6].extend(tool(suffile([cont_base, inst2, cont2])))
        message_list[6].extend(assistant(ans_o1))

        dataset_message_candidate.extend(message_list)

    return dataset_message_candidate


def build_output_dataset(
    datasets: List[Dict[str, Any]],
    dataset_message_candidate: List[List[Dict[str, str]]],
    rng: random.Random,
) -> List[Dict[str, Any]]:
    """
    output_dataset: per your pseudocode, for each record in datasets we build one temp_message_list by
    concatenating random(1,10) candidates; after each concat, remove last assistant(ans).
    We store each produced conversation as a JSONL row: {"messages": [...]}
    """
    output_rows: List[Dict[str, Any]] = []

    for candidate in dataset_message_candidate:
        temp_message_list: List[Dict[str, str]] = []

        # random(1,10) -> interpret as randint(1, 9) because Python range(1,10) yields 1..9
        for _i in range(rng.randrange(1, 2)):
            if _i == 0:
                cand = candidate
            else:
                cand = rng.choice(dataset_message_candidate)
            temp_message_list.extend(cand)

        # remove last assistant(ans)
        if temp_message_list and temp_message_list[-1].get("role") == "assistant":
            temp_message_list.pop()

        output_rows.append({"messages": temp_message_list})

    return output_rows


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--orig-inj-datapath", required=True, help="Path to original JSON dataset (list of dicts).")
    ap.add_argument("--output-path", required=True, help="Output JSONL path.")
    args = ap.parse_args()

    in_path = Path(args.orig_inj_datapath)
    out_path = Path(args.output_path)

    datasets = read_json(in_path)
    if not isinstance(datasets, list):
        raise ValueError("Input JSON must be a list of records (dict).")

    # Extract fields as you requested (even though later sampling uses 'other data')
    # This also sanity-checks schema early.
    for idx, d in enumerate(datasets):
        if not isinstance(d, dict):
            raise ValueError(f"Record {idx} is not a dict.")
        for key in ("instruction", "input", "output", "injection"):
            if key not in d:
                raise ValueError(f"Record {idx} missing required key: {key}")

    rng = random.Random(42)

    dataset_message_candidate = build_candidates(datasets, rng)
    print(len(dataset_message_candidate))
    output_rows = build_output_dataset(datasets, dataset_message_candidate, rng)

    write_jsonl(out_path, output_rows)
    print(f"Wrote {len(output_rows)} conversations to {out_path}")


if __name__ == "__main__":
    main()