from __future__ import annotations

import json
import os
from pathlib import Path

from datasets import load_dataset
from dotenv import load_dotenv


load_dotenv()

OUTPUT_DIR = Path("data/raw_sources/tensortrust")
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

OUTPUT_FILE = OUTPUT_DIR / "tensortrust_full.jsonl"
HF_TOKEN = os.getenv("HUGGINGFACE_TOKEN")


def load_tensortrust_split(file_path: str):
    kwargs = {
        "path": "json",
        "data_files": file_path,
        "split": "train",
    }

    if HF_TOKEN:
        kwargs["token"] = HF_TOKEN

    return load_dataset(**kwargs)


def main() -> None:
    hijacking_path = (
        "hf://datasets/qxcv/tensor-trust/benchmarks/"
        "hijacking-robustness/v1/hijacking_robustness_dataset.jsonl"
    )
    extraction_path = (
        "hf://datasets/qxcv/tensor-trust/benchmarks/"
        "extraction-robustness/v1/extraction_robustness_dataset.jsonl"
    )

    print("Downloading TensorTrust hijacking split...")
    hijacking_ds = load_tensortrust_split(hijacking_path)
    print(f"Loaded {len(hijacking_ds)} hijacking records.")

    print("Downloading TensorTrust extraction split...")
    extraction_ds = load_tensortrust_split(extraction_path)
    print(f"Loaded {len(extraction_ds)} extraction records.")

    with OUTPUT_FILE.open("w", encoding="utf-8") as f:
        for row in hijacking_ds:
            enriched = dict(row)
            enriched["benchmark_name"] = "hijacking"
            f.write(json.dumps(enriched, ensure_ascii=False) + "\n")
        for row in extraction_ds:
            enriched = dict(row)
            enriched["benchmark_name"] = "extraction"
            f.write(json.dumps(enriched, ensure_ascii=False) + "\n")

    print(f"Saved combined TensorTrust dataset to: {OUTPUT_FILE}")
    print(f"Total rows written: {len(hijacking_ds) + len(extraction_ds)}")


if __name__ == "__main__":
    main()