"""
consent_receipt.py — DPDPA Consent Management API Cryptographic Library
========================================================================

Production-grade Python 3.12+ module providing:
  - ConsentReceiptBuilder: Build and verify signed consent receipts (Section 6(10))
  - HashChainAuditLog: Compute and verify the append-only audit event hash chain
                       (Rule 6(1)(c), Rule 6(1)(e))
  - PseudonymService: HMAC-based principal pseudonymisation to prevent the audit
                      log from becoming a shadow PII store

Legal anchors:
  Section 5(3)    — notice must be available in 22 Eighth Schedule languages + English
  Section 6(1)    — consent must be free, specific, informed, unconditional, unambiguous
  Section 6(10)   — Data Fiduciary bears burden of proof; receipt is the evidence artefact
  Rule 6(1)(c)    — visibility on access to personal data via logs, monitoring, review
  Rule 6(1)(e)    — minimum 1-year log retention (see also Rule 8(3))

Dependencies:
  cryptography>=42.0   pip install cryptography
  (stdlib)             hashlib, hmac, json, datetime

Usage:
  from consent_receipt import ConsentReceiptBuilder, HashChainAuditLog, PseudonymService

  # Build a receipt
  builder = ConsentReceiptBuilder(signing_key=my_ed25519_private_key)
  receipt = builder.build(
      principal_id="user-uuid",
      purpose_id="marketing-purpose-uuid",
      notice_version_id="notice-v2-uuid",
      content_hash_shown=bytes.fromhex("a3f2c1d4..."),
      language_shown="ta",
      pseudonym_key=b"kms-held-hmac-key",
  )

  # Verify a receipt offline
  ok = ConsentReceiptBuilder.verify(receipt, public_key_pem=my_pub_key_pem)

  # Append an audit event
  log = HashChainAuditLog()
  event_hash = log.compute_event_hash(
      prev_hash=b"\\x00",
      consent_id="consent-uuid",
      event_type="GRANT",
      payload={"receipt_hash": receipt["receipt_hash"]},
      occurred_at=datetime.now(timezone.utc),
  )

  # Verify a chain
  is_valid, broken_links = log.verify_chain(events)
"""

from __future__ import annotations

import hashlib
import hmac
import json
from datetime import datetime, timezone
from typing import Any


# ---------------------------------------------------------------------------
# Optional import: cryptography library (required for signing/verification)
# ---------------------------------------------------------------------------
try:
    from cryptography.hazmat.primitives.asymmetric.ed25519 import (
        Ed25519PrivateKey,
        Ed25519PublicKey,
    )
    from cryptography.hazmat.primitives.serialization import (
        Encoding,
        PublicFormat,
        load_pem_public_key,
    )
    from cryptography.exceptions import InvalidSignature

    _CRYPTO_AVAILABLE = True
except ImportError:  # pragma: no cover
    _CRYPTO_AVAILABLE = False
    Ed25519PrivateKey = Any  # type: ignore[misc,assignment]
    Ed25519PublicKey = Any  # type: ignore[misc,assignment]


# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

#: Receipt format version string. Increment when the canonical format changes.
RECEIPT_VERSION = "DPDPA-CR-v1"

#: Separator used in the canonical receipt string.
_SEP = "|"

#: Sentinel prev_hash for the first event in a hash chain (single zero byte).
CHAIN_SENTINEL = b"\x00"

#: All 22 Eighth Schedule languages + English. Section 5(3).
EIGHTH_SCHEDULE_LANGUAGES: frozenset[str] = frozenset([
    "en",   # English
    "as",   # Assamese
    "bn",   # Bengali
    "brx",  # Bodo (ISO 639-3; BCP-47 approximation)
    "doi",  # Dogri
    "gu",   # Gujarati
    "hi",   # Hindi
    "kn",   # Kannada
    "ks",   # Kashmiri
    "kok",  # Konkani
    "mai",  # Maithili
    "ml",   # Malayalam
    "mni",  # Manipuri (Meitei)
    "mr",   # Marathi
    "ne",   # Nepali
    "or",   # Odia
    "pa",   # Punjabi
    "sa",   # Sanskrit
    "sat",  # Santali
    "sd",   # Sindhi
    "ta",   # Tamil
    "te",   # Telugu
    "ur",   # Urdu
])


# ---------------------------------------------------------------------------
# ConsentReceiptBuilder
# ---------------------------------------------------------------------------

class ConsentReceiptBuilder:
    """
    Build and verify Ed25519-signed consent receipts.

    A consent receipt is the primary Section 6(10) evidence artefact. It proves:
      1. A notice was given: notice_version_id + content_hash_shown identify the
         exact bytes of the notice that were rendered to the Data Principal.
      2. Consent was given: principal_pseudonym + granted_at timestamp record the act.
      3. Both were in accordance with the Act: the Ed25519 signature is verifiable
         offline against the published public key without accessing the API.

    # Section 6(10): burden of proof sits with the Data Fiduciary
    The receipt must be returned to the Data Principal at the time of grant.
    The Fiduciary retains only receipt_hash (SHA-256 of the canonical string)
    in the consent_record table — never the raw principal_id or full receipt.
    """

    def __init__(self, signing_key: Ed25519PrivateKey) -> None:
        """
        Initialise with an Ed25519 private key, typically loaded from KMS.

        Args:
            signing_key: Ed25519PrivateKey instance. In production, wrap an
                         AWS KMS / GCP Cloud KMS / Azure Key Vault signing call
                         rather than holding the raw key material in memory.
        """
        if not _CRYPTO_AVAILABLE:  # pragma: no cover
            raise ImportError(
                "The 'cryptography' package is required: pip install cryptography"
            )
        self._signing_key = signing_key

    # Section 6(10): build the evidence artefact
    def build(
        self,
        principal_id: str,
        purpose_id: str,
        notice_version_id: str,
        content_hash_shown: bytes,
        language_shown: str,
        pseudonym_key: bytes,
    ) -> dict[str, Any]:
        """
        Build a signed consent receipt for a single DPDPA consent grant.

        The canonical receipt string is:
          "DPDPA-CR-v1|{principal_pseudonym}|{purpose_id}|{notice_version_id}
           |{content_hash_hex}|{language_shown}|{granted_at_iso}"

        Separation with "|" makes the canonical form unambiguous and reproducible
        by any verifier, including an offline auditor with only the public key.

        Args:
            principal_id:       Raw internal principal UUID. NEVER stored as-is;
                                pseudonymised via HMAC before inclusion in the receipt.
            purpose_id:         UUID of the processing purpose. Section 6(1).
            notice_version_id:  UUID of the active notice version the client rendered.
                                Must match a non-superseded notice_version row.
            content_hash_shown: SHA-256 bytes of the notice content actually rendered
                                by the client. Verified server-side against
                                notice_translation.content_hash before this is called.
            language_shown:     BCP-47 code of the language the Data Principal chose.
                                Section 5(3).
            pseudonym_key:      HMAC-SHA256 key bytes, held in KMS. Annual rotation
                                controls relinkability. See PseudonymService.

        Returns:
            dict with fields: version, principal_pseudonym, purpose_id,
            notice_version_id, content_hash, language_shown, granted_at,
            receipt_hash, signature.

        Raises:
            ValueError: if language_shown is not in EIGHTH_SCHEDULE_LANGUAGES.
        """
        if language_shown not in EIGHTH_SCHEDULE_LANGUAGES:
            raise ValueError(
                f"language_shown '{language_shown}' is not a supported Eighth Schedule "
                f"language. Section 5(3) requires notice in one of: "
                f"{sorted(EIGHTH_SCHEDULE_LANGUAGES)}"
            )

        # Section 6(10): use pseudonym, not raw principal_id, in the receipt
        pseudonym_svc = PseudonymService(pseudonym_key)
        principal_pseudonym = pseudonym_svc.pseudonymise(principal_id)

        # Always timezone-aware. Never datetime.utcnow() — deprecated and naive.
        # A naive timestamp on a consent receipt is a defect an auditor will flag.
        granted_at: str = datetime.now(timezone.utc).isoformat()

        # Build canonical string — order is part of the protocol specification
        canonical_parts = [
            RECEIPT_VERSION,
            principal_pseudonym,
            purpose_id,
            notice_version_id,
            content_hash_shown.hex(),
            language_shown,
            granted_at,
        ]
        canonical: bytes = _SEP.join(canonical_parts).encode("utf-8")

        # Section 6(10): SHA-256 of canonical string, then Ed25519 sign
        receipt_hash: bytes = hashlib.sha256(canonical).digest()
        signature: bytes = self._signing_key.sign(receipt_hash)

        return {
            "version": RECEIPT_VERSION,
            "principal_pseudonym": principal_pseudonym,
            "purpose_id": purpose_id,
            "notice_version_id": notice_version_id,
            "content_hash": content_hash_shown.hex(),
            "language_shown": language_shown,
            "granted_at": granted_at,
            "receipt_hash": receipt_hash.hex(),
            "signature": signature.hex(),
        }

    @staticmethod
    def verify(receipt: dict[str, Any], public_key_pem: str) -> bool:
        """
        Offline receipt verifier. Can be run without database access.

        Reconstructs the canonical string from the receipt fields, re-computes
        receipt_hash, then verifies the Ed25519 signature against the provided
        public key.

        This is the operation an auditor or the Data Protection Board would run
        to verify a receipt without needing access to the Data Fiduciary's systems.
        Section 6(10): verifiability is a requirement of the burden of proof.

        Args:
            receipt:        dict as returned by ConsentReceiptBuilder.build().
            public_key_pem: PEM-encoded Ed25519 public key string. Should be the
                            Data Fiduciary's published signing key from
                            /.well-known/consent-signing-key.pub.

        Returns:
            True if the receipt is authentic and unmodified.
            False if the signature is invalid or any field has been tampered with.

        Note:
            This method does NOT validate that the notice_version_id or
            purpose_id correspond to known database records. Call the
            GET /v1/consents/{consent_id}/receipt endpoint for full validation.
        """
        if not _CRYPTO_AVAILABLE:  # pragma: no cover
            raise ImportError(
                "The 'cryptography' package is required: pip install cryptography"
            )

        try:
            canonical_parts = [
                receipt["version"],
                receipt["principal_pseudonym"],
                receipt["purpose_id"],
                receipt["notice_version_id"],
                receipt["content_hash"],
                receipt["language_shown"],
                receipt["granted_at"],
            ]
            canonical: bytes = _SEP.join(canonical_parts).encode("utf-8")
            receipt_hash: bytes = hashlib.sha256(canonical).digest()

            # Verify claimed receipt_hash matches recomputed value
            claimed_hash = bytes.fromhex(receipt["receipt_hash"])
            if not hmac.compare_digest(claimed_hash, receipt_hash):
                return False

            # Verify Ed25519 signature
            public_key: Ed25519PublicKey = load_pem_public_key(
                public_key_pem.encode("utf-8")
            )  # type: ignore[assignment]
            signature: bytes = bytes.fromhex(receipt["signature"])
            public_key.verify(signature, receipt_hash)
            return True

        except (InvalidSignature, KeyError, ValueError, TypeError):
            return False


# ---------------------------------------------------------------------------
# Payload type guard
# ---------------------------------------------------------------------------

_JSON_SAFE_INT_MAX = 2 ** 53
_JSON_SAFE_INT_MIN = -(2 ** 53)


def _assert_safe_payload(value: Any, path: str = "payload") -> None:
    """
    Recursively validate that a payload contains no float or unsafe int values.

    PostgreSQL jsonb normalizes numerics differently from Python json.dumps:
      - Floats with exponent notation: Python '1e+20' vs SQL '100000000000000000000'
      - Integers outside ±2^53: PostgreSQL may emit full integer form while Python
        preserves the original representation.

    Any divergence here causes the SQL and Python hash preimages to differ,
    breaking HashChainAuditLog.verify_chain() for affected events.

    Raises:
        ValueError: with the JSON path of the offending value.
    """
    if isinstance(value, bool):
        return  # bool is a subclass of int; check before int
    if isinstance(value, float):
        raise ValueError(
            f"{path}: float values are not allowed in audit payloads. "
            "PostgreSQL jsonb serializes floats differently from Python "
            "(e.g. 1e20 → '100000000000000000000' in SQL vs '1e+20' in Python). "
            "Convert to str before adding to the payload."
        )
    if isinstance(value, int) and not (_JSON_SAFE_INT_MIN <= value <= _JSON_SAFE_INT_MAX):
        raise ValueError(
            f"{path}: integer {value!r} is outside the safe JSON range "
            f"({_JSON_SAFE_INT_MIN} to {_JSON_SAFE_INT_MAX}). "
            "Large integers are normalized differently by PostgreSQL jsonb. "
            "Convert to str before adding to the payload."
        )
    if isinstance(value, dict):
        for k, v in value.items():
            _assert_safe_payload(v, path=f"{path}.{k}")
    elif isinstance(value, list):
        for i, v in enumerate(value):
            _assert_safe_payload(v, path=f"{path}[{i}]")


# ---------------------------------------------------------------------------
# HashChainAuditLog
# ---------------------------------------------------------------------------

class HashChainAuditLog:
    """
    Compute and verify the hash-chained consent audit event log.

    Each event in the log is linked to the previous one via:
      this_event_hash = sha256(
          prev_hash || consent_id || event_type || occurred_at_iso || payload_canonical
      )

    Any retroactive modification to a historical event produces a hash that
    diverges from all subsequent events, making tampering detectable by walking
    the chain forwards.

    # Rule 6(1)(c): every consent lifecycle event must be logged
    # Rule 6(1)(e): minimum 1-year retention; evidentiary retention per Section 6(10)
    # Section 6(10): tamper-evidence requires the hash chain to be anchored
    #                externally (e.g. timestamping authority, transparency log)

    Event types (as defined in consent_audit_event):
      NOTICE_SHOWN, GRANT, VIEW, WITHDRAW, ERASURE_FANOUT, ERASURE_CONFIRMED
    """

    #: Separator used in hash preimage construction.
    _SEP = b"|"

    def compute_event_hash(
        self,
        prev_hash: bytes,
        consent_id: str,
        event_type: str,
        payload: dict[str, Any],
        occurred_at: datetime,
    ) -> bytes:
        """
        Compute the hash chain entry for a single consent lifecycle event.

        Algorithm (matches append_audit_event() SQL function via jsonb_canonical()):
          payload_canonical = json.dumps(payload, sort_keys=True)
            # Pure-lex key ordering, default separators (", " and ": ").
            # Do NOT add separators=(',', ':') — that produces compact form which
            # diverges from PostgreSQL jsonb_canonical(). The two implementations
            # must agree on separators AND key ordering for verify_chain() to pass.
          preimage = prev_hash || '|' || consent_id_utf8 || '|' || event_type_utf8
                     || '|' || occurred_at_iso_utf8 || '|' || payload_canonical_utf8
          this_hash = sha256(preimage)

        Args:
            prev_hash:   bytes of previous event's this_event_hash, or CHAIN_SENTINEL
                         (b'\\x00') for the first event. Never pass an empty bytes.
            consent_id:  UUID string of the consent record.
            event_type:  One of: NOTICE_SHOWN, GRANT, VIEW, WITHDRAW,
                                 ERASURE_FANOUT, ERASURE_CONFIRMED.
            payload:     Event payload dict. Must not contain raw principal_id.
                         Keys and string values must be serialisable to JSON.
                         Do NOT include float values or integers outside the safe
                         JSON range (-(2**53) to 2**53): PostgreSQL jsonb normalizes
                         large ints and exponent floats differently from Python
                         (e.g. 1e20 → 100000000000000000000 in SQL vs '1e+20' in
                         Python), causing hash divergence. All payload values must
                         be str, bool, None, or nested dicts/lists of the same.
            occurred_at: timezone-aware datetime. Always use datetime.now(timezone.utc).
                         Never datetime.utcnow() — naive timestamps are not auditable.
                         # Section 6(10): timestamp is part of the verifiable evidence

        Returns:
            32-byte SHA-256 digest of the preimage.

        Raises:
            ValueError: if occurred_at is naive (no tzinfo), or if payload
                        contains float or out-of-safe-range int values.
        """
        if occurred_at.tzinfo is None:
            raise ValueError(
                "occurred_at must be a timezone-aware datetime. "
                "Use datetime.now(timezone.utc), not datetime.utcnow(). "
                "Naive timestamps are not auditable under Section 6(10)."
            )
        _assert_safe_payload(payload)

        # Canonicalise payload: sorted keys, standard separators matching jsonb_canonical().
        # ensure_ascii=False: PostgreSQL emits UTF-8 literals; Python's default ASCII-
        # escapes non-ASCII characters (e.g. "café" → "caf\u00e9"), causing hash
        # divergence for any non-ASCII key or string value. Do NOT remove ensure_ascii=False.
        # Do NOT add separators=(',', ':') — compact form diverges from jsonb_canonical().
        payload_canonical: bytes = json.dumps(
            payload, sort_keys=True, ensure_ascii=False
        ).encode("utf-8")

        # ISO 8601 with microseconds and UTC offset — matches PostgreSQL output
        occurred_at_iso: bytes = occurred_at.isoformat().encode("utf-8")

        preimage = (
            prev_hash
            + self._SEP
            + consent_id.encode("utf-8")
            + self._SEP
            + event_type.encode("utf-8")
            + self._SEP
            + occurred_at_iso
            + self._SEP
            + payload_canonical
        )
        return hashlib.sha256(preimage).digest()

    def verify_chain(
        self, events: list[dict[str, Any]]
    ) -> tuple[bool, list[str]]:
        """
        Walk the hash chain and verify every link is intact.

        Events must be provided in ascending event_id order (as returned by
        GET /v1/audit/{consent_id}). Each event dict must have:
          - event_id (int)
          - consent_id (str, UUID)
          - event_type (str)
          - prev_event_hash (str, hex — "00" or a 64-char hex string)
          - this_event_hash (str, hex — 64-char hex string)
          - event_payload (dict)
          - occurred_at (str, ISO 8601 with timezone)

        # Section 6(10): a broken chain link indicates tampering with audit evidence

        Args:
            events: List of audit event dicts in ascending event_id order.

        Returns:
            (True, []) if the chain is intact and every link verifies.
            (False, [list_of_broken_link_descriptions]) if any link fails.
        """
        if not events:
            return True, []

        broken_links: list[str] = []

        for i, event in enumerate(events):
            event_id = event.get("event_id", f"index_{i}")

            # Parse prev_hash
            prev_hex = event.get("prev_event_hash", "00")
            if prev_hex in ("00", None, ""):
                prev_hash = CHAIN_SENTINEL
            else:
                try:
                    prev_hash = bytes.fromhex(prev_hex)
                except ValueError:
                    broken_links.append(
                        f"event_id={event_id}: invalid hex in prev_event_hash '{prev_hex}'"
                    )
                    continue

            # Parse occurred_at — must be timezone-aware for deterministic replay
            try:
                occurred_at_raw = event["occurred_at"]
                occurred_at = datetime.fromisoformat(occurred_at_raw)
                if occurred_at.tzinfo is None:
                    # Assume UTC if no offset (legacy records only)
                    occurred_at = occurred_at.replace(tzinfo=timezone.utc)
            except (KeyError, ValueError) as exc:
                broken_links.append(
                    f"event_id={event_id}: unparseable occurred_at — {exc}"
                )
                continue

            # Recompute hash
            try:
                computed_hash = self.compute_event_hash(
                    prev_hash=prev_hash,
                    consent_id=event["consent_id"],
                    event_type=event["event_type"],
                    payload=event.get("event_payload", {}),
                    occurred_at=occurred_at,
                )
            except (KeyError, ValueError) as exc:
                broken_links.append(
                    f"event_id={event_id}: hash computation failed — {exc}"
                )
                continue

            # Compare against stored hash
            try:
                stored_hash = bytes.fromhex(event["this_event_hash"])
            except (KeyError, ValueError) as exc:
                broken_links.append(
                    f"event_id={event_id}: invalid hex in this_event_hash — {exc}"
                )
                continue

            if not hmac.compare_digest(computed_hash, stored_hash):
                broken_links.append(
                    f"event_id={event_id}: hash chain break — "
                    f"recomputed={computed_hash.hex()} stored={event['this_event_hash']}"
                )

        return (len(broken_links) == 0), broken_links


# ---------------------------------------------------------------------------
# PseudonymService
# ---------------------------------------------------------------------------

class PseudonymService:
    """
    HMAC-SHA256-based principal pseudonymisation.

    The consent audit log must not become a secondary PII store. Storing the
    raw principal_id in consent_audit_event.event_payload would mean the audit
    table is a directory of the user base, creating a Section 8(5) liability.

    The pseudonym approach:
      pseudonym = HMAC-SHA256(principal_id.encode('utf-8'), key=hmac_key)

    Properties:
    - Deterministic: the same (principal_id, key) always produces the same pseudonym,
      so audit rows can be aggregated for a given principal during investigation.
    - One-way: without the key, an attacker with the pseudonym cannot reverse it
      to find the principal_id (unlike a simple hash with a known format).
    - Rotatable: replacing hmac_key breaks the mapping from old pseudonyms to
      current principal_ids. Old audit rows remain self-consistent but cannot be
      bulk-linked to the current user base after rotation.

    # Annual key rotation for relinkability control
    The key should be stored in KMS and rotated annually. After rotation:
    - New consents use the new key (new pseudonym for each principal).
    - Old audit rows retain their old pseudonyms and remain internally consistent.
    - An attacker who exfiltrates old rows cannot link them to current users
      without both the old key AND the current principal_id mapping.
    - The Data Fiduciary can still look up a specific consent by principal_id
      for a rights request (Sections 11-12) by recomputing the pseudonym under the
      appropriate key for the time period.

    Rotation schedule: align with your annual KMS key rotation policy. (Rule 6(1)
    requires reasonable security safeguards but does not set a rotation period.)
    """

    def __init__(self, hmac_key: bytes) -> None:
        """
        Args:
            hmac_key: Raw bytes of the HMAC key from KMS. Must be at least 32 bytes.
                      In production, load from AWS KMS GenerateDataKey or equivalent.

        Raises:
            ValueError: if hmac_key is shorter than 32 bytes.
        """
        if len(hmac_key) < 32:
            raise ValueError(
                "hmac_key must be at least 32 bytes. "
                "Use a KMS-generated data key or equivalent."
            )
        self._hmac_key = hmac_key

    def pseudonymise(self, principal_id: str) -> str:
        """
        Compute the pseudonym for a given principal_id.

        Returns a 64-character lowercase hex string (SHA-256 output).

        Args:
            principal_id: Raw internal identifier (UUID string or similar).

        Returns:
            Hex-encoded HMAC-SHA256 digest.

        Usage:
            pseudonym = service.pseudonymise("d1e2f3a4-0000-0000-0000-000000000001")
            # → "e3b0c44298fc1c14..." (64 hex chars)
        """
        return hmac.new(
            self._hmac_key,
            principal_id.encode("utf-8"),
            hashlib.sha256,
        ).hexdigest()

    def pseudonymise_batch(self, principal_ids: list[str]) -> dict[str, str]:
        """
        Pseudonymise a list of principal_ids in one call.

        Useful for rights-request responses (Sections 11-12) where all consents for
        a principal need to be located across the audit log.

        Args:
            principal_ids: List of raw principal identifier strings.

        Returns:
            dict mapping principal_id → pseudonym.
        """
        return {pid: self.pseudonymise(pid) for pid in principal_ids}
