#!/usr/bin/env python3
"""Pinned official DataFog regex baseline on shared email/phone scope."""
import argparse
import hashlib
import importlib.metadata
import json
import sys
import time
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from pii.evaluation import evaluate, read_jsonl, write_json

VERSION = '4.8.1'
MAPPING = {'EMAIL': 'email', 'PHONE': 'phone_number'}


def scoped_gold(records):
    return [{**r, 'spans':[s for s in r['spans'] if s['label'] in MAPPING.values()]} for r in records]


def predict(records):
    import datafog
    if importlib.metadata.version('datafog') != VERSION:
        raise RuntimeError(f'Baseline requires datafog=={VERSION}')
    for r in records:
        result = datafog.scan(r['text'], engine='regex', entity_types=list(MAPPING))
        findings = []
        byte_offsets = [0]
        for ch in r['text']:
            byte_offsets.append(byte_offsets[-1] + len(ch.encode('utf-8')))
        for e in result.entities:
            if r['text'][e.start:e.end] != e.text:
                raise ValueError('Upstream span/source mismatch')
            findings.append(dict(label=MAPPING[e.type], start=byte_offsets[e.start], end=byte_offsets[e.end], raw_score=1.0, confidence=None))
        yield dict(id=r['id'], text=r['text'], findings=findings, contains_pii=bool(findings), score_semantics='Constant match indicator from DataFog regex, not calibrated confidence')


def write_jsonl(path, records):
    with open(path, 'w') as f:
        for record in records:
            f.write(json.dumps(record, ensure_ascii=False)+'\n')


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--gold', required=True)
    parser.add_argument('--output-dir', required=True)
    parser.add_argument('--model-predictions')
    parser.add_argument('--model-calibration')
    args = parser.parse_args()
    output = Path(args.output_dir)
    output.mkdir(parents=True, exist_ok=True)
    gold = scoped_gold(read_jsonl(args.gold))
    started = time.perf_counter()
    predictions = list(predict(gold))
    elapsed = time.perf_counter()-started
    write_jsonl(output/'gold-email-phone.jsonl', gold)
    write_jsonl(output/'regex-predictions.jsonl', predictions)
    metrics = evaluate(gold, predictions)
    # A regex match flag is not a probability; do not imply a calibration comparison.
    for key in ('raw_reliability','confidence_reliability','reliability_population','confidence_semantics'):
        metrics.pop(key, None)
    metrics['score_semantics'] = 'Regex match indicator, no probability calibration evaluated'
    metrics['comparison_scope'] = ['email','phone_number']
    metrics['unsupported_labels'] = ['first_name','last_name','street_address','customer_id','account_number']
    write_json(output/'regex-metrics.json', metrics)
    provenance = dict(package='datafog', version=VERSION, official_repository='https://github.com/DataFog/datafog-python',
                      release='https://pypi.org/project/datafog/4.8.1/', input_sha256=hashlib.sha256(Path(args.gold).read_bytes()).hexdigest(),
                      documents=len(gold), elapsed_seconds=elapsed, API="datafog.scan(text, engine='regex', entity_types=['EMAIL','PHONE'])",
                      offset_conversion='Python character offsets converted to original UTF-8 byte offsets', dependencies={d.metadata['Name']:d.version for d in importlib.metadata.distributions()})
    write_json(output/'regex-provenance.json', provenance)
    if args.model_predictions:
        model = read_jsonl(args.model_predictions)
        model = [{**r, 'findings':[s for s in r['findings'] if s['label'] in MAPPING.values()]} for r in model]
        calibration = json.loads(Path(args.model_calibration).read_text()) if args.model_calibration else None
        model_metrics = evaluate(gold, model, calibration=calibration)
        model_metrics['comparison_scope'] = ['email','phone_number']
        write_json(output/'model-email-phone-metrics.json', model_metrics)
    print(json.dumps({'documents':len(gold), 'seconds':elapsed, 'output_dir':str(output)}))