#!/usr/bin/env python3
"""Protect (encrypt) a PDF with a password using pikepdf."""
import argparse, json, sys
from pathlib import Path

def protect(input_path, output_path, password):
    import pikepdf
    
    if not output_path:
        output_path = str(Path(input_path).with_stem(Path(input_path).stem + '_protected'))
    
    pdf = pikepdf.open(input_path)
    pdf.save(output_path, encryption=pikepdf.Encryption(
        owner=password,
        user=password,
        aes=True,
        R=6  # AES-256
    ))
    pdf.close()
    return output_path

def main():
    parser = argparse.ArgumentParser(description='Protect PDF with password')
    parser.add_argument('--input', required=True)
    parser.add_argument('--output', required=True)
    parser.add_argument('--password', required=True)
    args = parser.parse_args()
    try:
        result = protect(args.input, args.output, args.password)
        print(json.dumps({"success": True, "output": result, "message": "PDF protected successfully"}))
    except Exception as e:
        print(json.dumps({"success": False, "output": "", "message": str(e)}))
        sys.exit(1)

if __name__ == '__main__':
    main()