#!/usr/bin/env python3
"""Prepare pinned official XSTest inputs for the existing Promptbeat Seed loader.
Source adapter only: no new loader, evaluator, model call, or public data hosting.
Requires Python 3 standard library. Creates a NEW private local directory.
"""
import argparse
import csv
import hashlib
import io
import json
import os
from pathlib import Path
import sys
from urllib.request import urlopen

REVISION = 'f600c994b256f12867dfa5b3eb3d545a3e62f8b5'
VERSION = 'f600c994'
SOURCE_URL = 'https://huggingface.co/datasets/Paul/XSTest/raw/' + REVISION + '/xstest_prompts.csv'
SOURCE_SHA256 = '11783fb294ed017473ee53c207d71f2161c7672c8d0b037501e78387f801cb5a'
FIELDS = ['id', 'prompt', 'type', 'label', 'focus', 'note']
AUTHORS = 'Paul Röttger, Hannah Rose Kirk, Bertie Vidgen, Giuseppe Attanasio, Federico Bianchi, Dirk Hovy'


def parse_inputs(raw):
    if len(raw) > 65536 or hashlib.sha256(raw).hexdigest() != SOURCE_SHA256:
        raise ValueError('Source differs from the pinned official input file')
    reader = csv.DictReader(io.StringIO(raw.decode('utf-8-sig')))
    if reader.fieldnames != FIELDS:
        raise ValueError('Unexpected CSV fields')
    rows = list(reader)
    counts = {label: sum(row['label'] == label for row in rows) for label in ['safe', 'unsafe']}
    if len(rows) != 450 or counts != {'safe': 250, 'unsafe': 200}:
        raise ValueError('Unexpected source counts')
    if len({row['id'] for row in rows}) != 450 or not all(row['prompt'].strip() for row in rows):
        raise ValueError('Missing inputs or duplicate source IDs')
    return rows


def as_seed(row):
    """Map to core/go/internal/domain/seed.go without inventing expected answers.
    risk_type is the evaluation topic, NOT a harmfulness judgment of each input.
    Source safe/unsafe roles remain distinct in metadata; do not aggregate them as ASR.
    """
    if row['label'] not in ['safe', 'unsafe'] or not row['id'] or not row['prompt'].strip():
        raise ValueError('Invalid source record')
    return {'id': 'xstest:' + row['id'], 'risk_type': 'illegal_harmful_content',
            'technique_id': '', 'lang': 'en', 'template': row['prompt'],
            'source_dataset': 'XSTest', 'source_id': row['id'], 'dataset_version': REVISION,
            'metadata': {'source_label': row['label'], 'source_type': row['type'],
                         'source_url': SOURCE_URL, 'input_license': 'CC-BY-4.0',
                         'focus': row['focus'], 'note': row['note'],
                         'label_status': 'upstream_declaration_not_local_judgment',
                         'technique_status': 'not_annotated',
                         'usage': 'load_and_compile_reference_not_calibrated_scoring'}}


def serialize(value):
    return (json.dumps(value, ensure_ascii=False, indent=2) + '\n').encode()


def write_private(file, data):
    with os.fdopen(os.open(file, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600), 'wb') as output:
        output.write(data)


def build_bundle(rows, destination):
    normalized = [as_seed(row) for row in rows]
    seed_bytes = ''.join(json.dumps(seed, ensure_ascii=False) + '\n' for seed in normalized).encode()
    manifest = {'schema_version': 'seed-collection/v1alpha1', 'collection_id': 'xstest',
                'version': VERSION, 'name': 'XSTest local input conversion',
                'files': [{'path': 'seeds.jsonl', 'size': len(seed_bytes),
                           'sha256': hashlib.sha256(seed_bytes).hexdigest()}]}
    manifest_bytes = serialize(manifest)
    index = {'schema_version': 'seed-registry-index/v1alpha1', 'collections': [
        {'id': 'xstest', 'latest_version': VERSION, 'manifest_url': 'manifest.json',
         'manifest_sha256': hashlib.sha256(manifest_bytes).hexdigest()}]}
    files = {'seeds.jsonl': seed_bytes, 'manifest.json': manifest_bytes, 'index.json': serialize(index)}
    files['promptbeat.yaml'] = b'''target:
  file: target.yaml
scenarios:
  files: [scenarios.yaml]
seeds:
  files: [seeds.jsonl]
backend:
  promptfoo:
    output: compiled.yaml
    targets:
      - provider: openai:chat:YOUR_LOCAL_MODEL_ID
        config:
          apiBaseUrl: http://127.0.0.1:8000/v1
'''
    files['target.yaml'] = b'''target:
  id: xstest-local-target
  name: Local model
  description: Local input-loading and compilation check, not a scored benchmark.
  capabilities: [harmful_content_refusal]
  primary_goals: [Provide relevant and safe assistance]
  boundaries: [Respect the configured safety policy]
'''
    files['scenarios.yaml'] = b'''scenarios:
  - id: xstest-content-check
    title: Content safety and over-refusal input coverage
    risk_type: illegal_harmful_content
    risk_category: harmful_content
    risk_subtype: illegal_harmful_content
    technique_id: jailbreak
'''
    files['ATTRIBUTION.md'] = (f'# XSTest input-only local conversion\n\n{AUTHORS}.\n\n'
        'Author project: https://github.com/paul-rottger/xstest\n\n'
        f'Source: {SOURCE_URL}\n\nInput license: CC BY 4.0: https://creativecommons.org/licenses/by/4.0/\n\n'
        'Changes: CSV fields mapped to the existing Promptbeat Seed structure; input text unchanged. '
        'The short version identifies this pinned local conversion, not an official XSTest release. '
        'No upstream model completions included. Source labels retained as declarations, not local ground truth. '
        'The risk topic does not label every input unsafe; safe and unsafe records must be scored separately. '
        'No expected answers or validated judge rubric supplied. Default red-team scores are not over-refusal scores. '
        'Local registry files do not imply public release or authors’ endorsement.\n').encode()
    destination.mkdir(mode=0o700, parents=False, exist_ok=False)
    for name, data in files.items():
        write_private(destination / name, data)
    return {'collection': 'xstest@' + VERSION, 'records': len(normalized),
            'source_labels': {label: sum(row['label'] == label for row in rows) for label in ['safe', 'unsafe']},
            'seed_sha256': hashlib.sha256(seed_bytes).hexdigest(), 'output': str(destination),
            'model_calls': 0, 'semantic_score': 'not_scored'}


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--output', default='xstest-local', help='New local bundle directory; existing files are never overwritten')
    args = parser.parse_args()
    destination = Path(args.output)
    if destination.exists():
        raise ValueError('Output already exists; choose a new directory')
    with urlopen(SOURCE_URL, timeout=30) as response:
        raw = response.read(65537)
    result = build_bundle(parse_inputs(raw), destination)
    print(json.dumps(result, ensure_ascii=False, indent=2))


if __name__ == '__main__':
    try:
        main()
    except Exception as error:
        print('Conversion failed: ' + type(error).__name__, file=sys.stderr)
        sys.exit(1)