#!/usr/bin/env python3
from __future__ import annotations

import json
import re
from pathlib import Path

from prompt_injection_framework.io.data_loader import read_jsonl_records, validate_core_fields


CURATED_SOURCES = (
    Path("data/curated/hackaprompt_successful_attacks.jsonl"),
    Path("data/curated/tensortrust_curated_attacks.jsonl"),
)
OUTPUT_PATH = Path("data/normalized/attack_cases.jsonl")

EXPECTED_COMPLETION_RE = re.compile(r"expected_completion=([^,\n]+)")


def hydrate_expected_target(case: dict) -> dict:
    hydrated = dict(case)
    if str(hydrated.get("expected_target", "")).strip():
        return hydrated

    notes = str(hydrated.get("notes", ""))
    match = EXPECTED_COMPLETION_RE.search(notes)
    if match:
        hydrated["expected_target"] = match.group(1).strip()
    return hydrated


def main() -> None:
    OUTPUT_PATH.parent.mkdir(parents=True, exist_ok=True)

    total = 0
    with OUTPUT_PATH.open("w", encoding="utf-8") as output_handle:
        for source_path in CURATED_SOURCES:
            for case in read_jsonl_records(source_path):
                hydrated_case = hydrate_expected_target(case)
                validate_core_fields(hydrated_case)
                output_handle.write(json.dumps(hydrated_case, ensure_ascii=False) + "\n")
                total += 1

    print(f"Wrote {total} normalized records to {OUTPUT_PATH}")


if __name__ == "__main__":
    main()