#!/usr/bin/env python3
"""
embed-c2pa.py
=============
Embeds a C2PA (Coalition for Content Provenance and Authenticity) manifest
into an AI-generated marketing asset (image / video / audio / PDF) so the
asset carries a machine-readable provenance trail and a visible AI-generation
claim. This is the technical mechanism for EU AI Act Article 50 compliance
(applicable 2 Aug 2026) and aligns with the broader Content Authenticity
Initiative (CAI) standard.

Usage:
    python3 embed-c2pa.py \
        --input path/to/asset.png \
        --output path/to/signed-asset.png \
        --brand "Acme Corp" \
        --generator "Vertex AI / Nano Banana Pro" \
        --ai-claim ai-generated-content \
        [--created 2026-05-16T10:30:00Z] \
        [--prompt "the prompt used"] \
        [--reviewer "Jane Smith"] \
        [--signing-cert /path/to/cert.pem --signing-key /path/to/key.pem]

If c2pa-python is not installed, it is auto-installed via pip. If a real
signing certificate is not provided, the script falls back to a self-signed
test certificate (for development only — production use must supply a real
certificate from a CAI-recognized authority).

AI claim values (mapped to C2paDigitalSourceType):
    ai-generated-content        # Asset fully generated by AI (TRAINED_ALGORITHMIC_MEDIA)
    ai-assisted-edits           # Human-created + AI-edited (COMPOSITE_WITH_TRAINED_ALGORITHMIC_MEDIA)
    ai-no-substantive-changes   # AI used but no semantic change (HUMAN_EDITS)

The resulting asset can be inspected at https://contentcredentials.org/verify
or with any C2PA-aware viewer (Photoshop, Lightroom, Truepic, etc.).

Tested with c2pa-python 0.32.6. The library API changed in 0.13+ — this script
uses the current `Builder` + `Signer.from_info(C2paSignerInfo)` pattern.

Exit codes:
    0  success
    2  c2pa-python install failed
    3  unsupported asset format
    4  signing failure
    5  invalid arguments
"""

import argparse
import json
import subprocess
import sys
from datetime import datetime, timezone
from pathlib import Path


SUPPORTED_FORMATS = {
    ".png": "image/png",
    ".jpg": "image/jpeg",
    ".jpeg": "image/jpeg",
    ".webp": "image/webp",
    ".gif": "image/gif",
    ".tiff": "image/tiff",
    ".mp4": "video/mp4",
    ".mov": "video/quicktime",
    ".webm": "video/webm",
    ".mp3": "audio/mpeg",
    ".wav": "audio/wav",
    ".pdf": "application/pdf",
}

# Map our marketing-friendly claim values to the c2pa-python library's
# C2paDigitalSourceType enum (which is the official IPTC-aligned vocabulary
# the EU AI Act enforcement tooling reads).
AI_CLAIM_TO_C2PA_TYPE = {
    "ai-generated-content":      "TRAINED_ALGORITHMIC_MEDIA",
    "ai-assisted-edits":         "COMPOSITE_WITH_TRAINED_ALGORITHMIC_MEDIA",
    "ai-no-substantive-changes": "HUMAN_EDITS",
}


def ensure_c2pa():
    """Install c2pa-python if missing. Returns the module."""
    try:
        import c2pa
        return c2pa
    except ImportError:
        pass
    print("Installing c2pa-python...", file=sys.stderr)
    try:
        subprocess.check_call(
            [sys.executable, "-m", "pip", "install", "--quiet", "c2pa-python>=0.32"]
        )
        import c2pa
        return c2pa
    except Exception as exc:
        print(f"ERROR: could not install c2pa-python: {exc}", file=sys.stderr)
        print(
            "Install manually: pip install c2pa-python>=0.32\n"
            "Docs: https://opensource.contentauthenticity.org/docs/python/",
            file=sys.stderr,
        )
        sys.exit(2)


def build_manifest_json(brand, generator, ai_claim, created, prompt, reviewer, asset_format):
    """Build a manifest JSON definition the c2pa-python Builder accepts.

    Note: digital_source_type is set via builder.set_intent(), not embedded
    in the actions assertion — that's the current c2pa-python pattern.
    """
    actions = [
        {
            "action": "c2pa.created",
            "when": created,
            "softwareAgent": {
                "name": generator,
                "version": "1.0",
            },
        }
    ]
    if prompt:
        actions.append({
            "action": "c2pa.opened",
            "when": created,
            "parameters": {"description": f"Source prompt: {prompt}"},
        })
    if reviewer:
        actions.append({
            "action": "c2pa.edited",
            "when": datetime.now(timezone.utc).isoformat(),
            "parameters": {"description": f"Human-reviewed by: {reviewer}"},
        })

    manifest = {
        "claim_generator_info": [{
            "name": "DigitalMarketingPro",
            "version": "3.4.1",
        }],
        "title": f"{brand} — AI-generated marketing asset",
        "format": SUPPORTED_FORMATS[asset_format],
        "assertions": [
            {
                "label": "c2pa.actions.v2",
                "data": {"actions": actions},
            },
            {
                "label": "stds.schema-org.CreativeWork",
                "data": {
                    "@context": "https://schema.org",
                    "@type": "CreativeWork",
                    "author": [{"@type": "Organization", "name": brand}],
                    "dateCreated": created,
                },
            },
        ],
    }
    return manifest


def generate_self_signed_cert(tmp_dir):
    """Generate a C2PA-conformant self-signed cert for DEV use only.

    C2PA enforces certificate rules per the spec (https://c2pa.org/specifications/
    specifications/2.0/specs/C2PA_Specification.html#x_509_certificates):
      - Extended Key Usage must include emailProtection (1.3.6.1.5.5.7.3.4)
        OR documentSigning (1.3.6.1.5.5.7.3.36) OR the C2PA-specific OID
      - Key Usage must include digitalSignature
      - Subject Key Identifier extension is required (for non-root certs;
        we set it on self-signed too for trust-list cleanliness)
      - basicConstraints CA=false for an end-entity cert
    A vanilla self-signed cert without these gets "the certificate is invalid".
    """
    from cryptography import x509
    from cryptography.x509.oid import NameOID, ExtendedKeyUsageOID
    from cryptography.hazmat.primitives import hashes, serialization
    from cryptography.hazmat.primitives.asymmetric import ec
    from datetime import timedelta

    key = ec.generate_private_key(ec.SECP256R1())
    public_key = key.public_key()

    subject = issuer = x509.Name([
        x509.NameAttribute(NameOID.COMMON_NAME, "DMP Dev Self-Signed C2PA"),
        x509.NameAttribute(NameOID.ORGANIZATION_NAME, "Digital Marketing Pro"),
        x509.NameAttribute(NameOID.COUNTRY_NAME, "US"),
    ])

    builder = (
        x509.CertificateBuilder()
        .subject_name(subject)
        .issuer_name(issuer)
        .public_key(public_key)
        .serial_number(x509.random_serial_number())
        .not_valid_before(datetime.now(timezone.utc) - timedelta(minutes=1))
        .not_valid_after(datetime.now(timezone.utc) + timedelta(days=90))
        .add_extension(
            x509.BasicConstraints(ca=False, path_length=None),
            critical=True,
        )
        .add_extension(
            x509.KeyUsage(
                digital_signature=True,
                content_commitment=False,
                key_encipherment=False,
                data_encipherment=False,
                key_agreement=False,
                key_cert_sign=False,
                crl_sign=False,
                encipher_only=False,
                decipher_only=False,
            ),
            critical=True,
        )
        .add_extension(
            x509.ExtendedKeyUsage([ExtendedKeyUsageOID.EMAIL_PROTECTION]),
            critical=False,
        )
        .add_extension(
            x509.SubjectKeyIdentifier.from_public_key(public_key),
            critical=False,
        )
        .add_extension(
            x509.AuthorityKeyIdentifier.from_issuer_public_key(public_key),
            critical=False,
        )
    )
    cert = builder.sign(key, hashes.SHA256())

    cert_path = Path(tmp_dir) / "dev-c2pa-cert.pem"
    key_path = Path(tmp_dir) / "dev-c2pa-key.pem"
    cert_path.write_bytes(cert.public_bytes(serialization.Encoding.PEM))
    key_path.write_bytes(
        key.private_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PrivateFormat.PKCS8,
            encryption_algorithm=serialization.NoEncryption(),
        )
    )
    return str(cert_path), str(key_path)


def main():
    parser = argparse.ArgumentParser(description="Embed a C2PA provenance manifest in an AI-generated marketing asset.")
    parser.add_argument("--input", required=True, help="Source asset path")
    parser.add_argument("--output", required=True, help="Output asset path (with manifest embedded)")
    parser.add_argument("--brand", required=True, help="Brand name (becomes CreativeWork.author)")
    parser.add_argument("--generator", required=True, help='Generator name, e.g. "Vertex AI / Nano Banana Pro" or "Runway Gen-4"')
    parser.add_argument(
        "--ai-claim",
        default="ai-generated-content",
        choices=sorted(AI_CLAIM_TO_C2PA_TYPE.keys()),
        help="Maps to a C2paDigitalSourceType enum value used by Builder.set_intent",
    )
    parser.add_argument("--created", help="ISO-8601 creation timestamp (default: now)")
    parser.add_argument("--prompt", help="Original generation prompt (optional, embedded in actions)")
    parser.add_argument("--reviewer", help="Human reviewer name (optional, embedded in actions)")
    parser.add_argument("--signing-cert", help="PEM signing certificate; if omitted, a 90-day self-signed dev cert is generated")
    parser.add_argument("--signing-key", help="PEM signing key (must accompany --signing-cert)")
    args = parser.parse_args()

    in_path = Path(args.input)
    out_path = Path(args.output)

    if not in_path.exists():
        print(f"ERROR: input not found: {in_path}", file=sys.stderr)
        sys.exit(5)

    ext = in_path.suffix.lower()
    if ext not in SUPPORTED_FORMATS:
        print(f"ERROR: unsupported format {ext}. Supported: {sorted(SUPPORTED_FORMATS)}", file=sys.stderr)
        sys.exit(3)

    created = args.created or datetime.now(timezone.utc).isoformat()
    manifest = build_manifest_json(args.brand, args.generator, args.ai_claim, created, args.prompt, args.reviewer, ext)

    c2pa = ensure_c2pa()

    cert_path, key_path = args.signing_cert, args.signing_key
    using_dev_cert = False
    if not (cert_path and key_path):
        import tempfile
        tmpdir = tempfile.mkdtemp(prefix="dmp-c2pa-")
        try:
            cert_path, key_path = generate_self_signed_cert(tmpdir)
            using_dev_cert = True
        except ImportError:
            print("ERROR: cryptography package required for self-signed cert. Install: pip install cryptography", file=sys.stderr)
            sys.exit(4)

    out_path.parent.mkdir(parents=True, exist_ok=True)

    try:
        cert_bytes = open(cert_path, "rb").read()
        key_bytes = open(key_path, "rb").read()

        # Real c2pa-python 0.32+ API
        signer_info = c2pa.C2paSignerInfo(
            alg=b"es256",
            sign_cert=cert_bytes,
            private_key=key_bytes,
            ta_url=b"http://timestamp.digicert.com",
        )
        signer = c2pa.Signer.from_info(signer_info)

        builder = c2pa.Builder(manifest)

        # Bind the digital-source-type claim via set_intent(CREATE, <type>)
        ds_type_name = AI_CLAIM_TO_C2PA_TYPE[args.ai_claim]
        ds_type = getattr(c2pa.C2paDigitalSourceType, ds_type_name)
        try:
            builder.set_intent(c2pa.C2paBuilderIntent.CREATE, ds_type)
        except Exception as exc:
            # set_intent is optional — manifest's assertions already carry the claim
            print(f"NOTE: set_intent failed ({exc}); proceeding with assertions-only claim", file=sys.stderr)

        # Sign the file
        builder.sign_file(str(in_path), str(out_path), signer=signer)
    except Exception as exc:
        print(f"ERROR: C2PA signing failed: {type(exc).__name__}: {exc}", file=sys.stderr)
        sys.exit(4)

    # Verify the manifest read back (round-trip integrity check)
    manifest_present = False
    manifest_summary = None
    try:
        mime = SUPPORTED_FORMATS[ext]
        with open(out_path, "rb") as fh:
            with c2pa.Reader(mime, fh) as reader:
                manifest_summary = json.loads(reader.json())
                manifest_present = bool(manifest_summary.get("active_manifest"))
    except Exception as exc:
        print(f"NOTE: could not read back manifest for verification ({type(exc).__name__}: {exc})", file=sys.stderr)

    result = {
        "status": "success",
        "input": str(in_path),
        "output": str(out_path),
        "size_bytes": out_path.stat().st_size,
        "brand": args.brand,
        "generator": args.generator,
        "ai_claim": args.ai_claim,
        "c2pa_digital_source_type": ds_type_name,
        "created": created,
        "manifest_assertions": [a["label"] for a in manifest["assertions"]],
        "manifest_embedded_and_verified": manifest_present,
        "using_dev_cert": using_dev_cert,
        "verify_url": "https://contentcredentials.org/verify",
    }
    if manifest_summary and "active_manifest" in manifest_summary:
        result["c2pa_active_manifest_id"] = manifest_summary["active_manifest"]
    if using_dev_cert:
        result["warning"] = (
            "Signed with self-signed dev cert (90-day validity). "
            "PRODUCTION USE REQUIRES a real signing certificate from a CAI-recognized authority. "
            "See https://opensource.contentauthenticity.org/docs/manifest/signing-manifests/"
        )
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()
