"""RFC 3161 Timestamp Utilities for Receipts

Provides helper functions to add RFC 3161 timestamps to any receipt type.
Integrates with tools/crypto/rfc3161_client.py for TSA operations.

Usage:
    from receipts.rfc3161_utils import add_rfc3161_timestamp

    receipt = {
        "schema": "receipt.merge.v1",
        "receipt_id": "merge-123",
        ...
    }

    # Add RFC 3161 timestamp (optional - graceful fallback)
    receipt = add_rfc3161_timestamp(receipt, use_stub=True)

    # Receipt now has rfc3161_tst field
    assert "rfc3161_tst" in receipt
"""

import hashlib
import json
import logging
from typing import Dict, Any, Optional

from tools.crypto.rfc3161_client import (
    get_timestamp,
    get_timestamp_stub,
    RFC3161Error,
)

logger = logging.getLogger(__name__)


def add_rfc3161_timestamp(
    receipt: Dict[str, Any],
    *,
    use_stub: bool = True,
    tsa_server: Optional[str] = None,
    hash_algorithm: str = "sha256",
    required: bool = False,
) -> Dict[str, Any]:
    """
    Add RFC 3161 timestamp to a receipt.

    Computes SHA-256 hash of receipt (excluding signature) and requests
    timestamp from TSA server. If TSA server is unavailable, uses stub
    (unless required=True).

    Args:
        receipt: Receipt dictionary (modified in-place)
        use_stub: Use stub implementation if TSA unavailable (default: True)
        tsa_server: TSA server URL (None = use default)
        hash_algorithm: Hash algorithm (sha256, sha384, sha512)
        required: Raise error if TSA unavailable (default: False)

    Returns:
        Modified receipt with rfc3161_tst field

    Raises:
        RFC3161Error: If required=True and TSA unavailable
    """
    # Exclude fields that vary by timestamp
    exclude_fields = {
        "signature",
        "timestamp",
        "rfc3161_tst",
        "transparency",
        "legal_notice",
        "custodian_attestation",
    }

    # Compute content hash
    content = {k: v for k, v in receipt.items() if k not in exclude_fields}
    content_bytes = json.dumps(content, sort_keys=True).encode("utf-8")

    # Use specified hash algorithm
    if hash_algorithm == "sha256":
        data_hash = hashlib.sha256(content_bytes).digest()
    elif hash_algorithm == "sha384":
        data_hash = hashlib.sha384(content_bytes).digest()
    elif hash_algorithm == "sha512":
        data_hash = hashlib.sha512(content_bytes).digest()
    else:
        raise ValueError(f"Unsupported hash algorithm: {hash_algorithm}")

    # Request timestamp
    try:
        if use_stub:
            # Use stub for testing/development
            tst_data = get_timestamp_stub(data_hash, hash_algorithm=hash_algorithm)
            logger.debug("Using RFC 3161 stub timestamp")
        else:
            # Use real TSA server
            tst_data = get_timestamp(
                data_hash,
                hash_algorithm=hash_algorithm,
                tsa_server=tsa_server,
                timeout=10,
                retry_count=2
            )
            logger.info(f"Got RFC 3161 timestamp from {tst_data['tsa_server']}")

        # Add to receipt
        receipt["rfc3161_tst"] = {
            "tsa_server": tst_data["tsa_server"],
            "tst_der_b64": tst_data["tst_der_b64"],
            "tst_hash_algorithm": hash_algorithm,
            "tst_timestamp": tst_data["tst_timestamp"]
        }

    except RFC3161Error as e:
        if required:
            raise
        else:
            logger.warning(f"RFC 3161 timestamp unavailable: {e}")
            # Continue without timestamp

    return receipt


def add_rfc3161_to_batch(
    receipts: list[Dict[str, Any]],
    *,
    use_stub: bool = True,
    tsa_server: Optional[str] = None,
    hash_algorithm: str = "sha256",
) -> list[Dict[str, Any]]:
    """
    Add RFC 3161 timestamps to multiple receipts.

    Args:
        receipts: List of receipt dictionaries
        use_stub: Use stub implementation if TSA unavailable
        tsa_server: TSA server URL
        hash_algorithm: Hash algorithm

    Returns:
        List of receipts with rfc3161_tst fields
    """
    for receipt in receipts:
        add_rfc3161_timestamp(
            receipt,
            use_stub=use_stub,
            tsa_server=tsa_server,
            hash_algorithm=hash_algorithm,
            required=False
        )

    return receipts


def verify_rfc3161_timestamp(
    receipt: Dict[str, Any],
    *,
    hash_algorithm: str = "sha256",
) -> bool:
    """
    Verify RFC 3161 timestamp in a receipt.

    Args:
        receipt: Receipt with rfc3161_tst field
        hash_algorithm: Hash algorithm used

    Returns:
        True if timestamp is valid

    Raises:
        ValueError: Receipt missing rfc3161_tst field
        RFC3161Error: Verification failed
    """
    if "rfc3161_tst" not in receipt:
        raise ValueError("Receipt missing rfc3161_tst field")

    from tools.crypto.rfc3161_client import verify_timestamp

    # Recompute content hash
    exclude_fields = {
        "signature",
        "timestamp",
        "rfc3161_tst",
        "transparency",
        "legal_notice",
        "custodian_attestation",
    }

    content = {k: v for k, v in receipt.items() if k not in exclude_fields}
    content_bytes = json.dumps(content, sort_keys=True).encode("utf-8")

    if hash_algorithm == "sha256":
        data_hash = hashlib.sha256(content_bytes).digest()
    elif hash_algorithm == "sha384":
        data_hash = hashlib.sha384(content_bytes).digest()
    elif hash_algorithm == "sha512":
        data_hash = hashlib.sha512(content_bytes).digest()
    else:
        raise ValueError(f"Unsupported hash algorithm: {hash_algorithm}")

    # Verify
    tst_data = receipt["rfc3161_tst"]
    return verify_timestamp(
        tst_data["tst_der_b64"],
        data_hash,
        hash_algorithm=hash_algorithm
    )


def enforce_rfc3161_timestamp(
    receipt: Dict[str, Any],
    *,
    hash_algorithm: str = "sha256",
) -> None:
    """
    Guardian-style guard: refuse processing if TSA verification fails.

    Raises RFC3161Error when the timestamp is missing or invalid. Use this in
    any path that depends on court-grade time so receipts cannot proceed with
    untrusted TSRs.
    """
    try:
        ok = verify_rfc3161_timestamp(receipt, hash_algorithm=hash_algorithm)
    except Exception as exc:
        # Propagate existing RFC3161Error/ValueError so callers can block.
        raise

    if not ok:
        raise RFC3161Error("RFC 3161 timestamp verification failed")


if __name__ == "__main__":
    # Demo
    import time

    print("RFC 3161 Timestamp Utilities Demo")
    print("=" * 70)

    # Sample receipt
    receipt = {
        "schema": "receipt.merge.v1",
        "receipt_id": "merge-demo-123",
        "agent_id": "agent-001",
        "file_path": "src/components/App.tsx",
        "hunks_merged": [
            {"claim_id": "claim-001", "start_line": 10, "end_line": 20}
        ],
        "guardian_verdict": "pass",
        "status": "applied",
        "timestamp": time.time()
    }

    print("\nOriginal receipt:")
    print(f"  receipt_id: {receipt['receipt_id']}")
    print(f"  has rfc3161_tst: {'rfc3161_tst' in receipt}")

    # Add timestamp
    print("\nAdding RFC 3161 timestamp...")
    receipt = add_rfc3161_timestamp(receipt, use_stub=True)

    print(f"  has rfc3161_tst: {'rfc3161_tst' in receipt}")
    if "rfc3161_tst" in receipt:
        print(f"  tsa_server: {receipt['rfc3161_tst']['tsa_server']}")
        print(f"  tst_timestamp: {receipt['rfc3161_tst']['tst_timestamp']}")
        print(f"  tst_der_b64: {receipt['rfc3161_tst']['tst_der_b64'][:32]}...")

    # Verify
    print("\nVerifying timestamp...")
    try:
        is_valid = verify_rfc3161_timestamp(receipt)
        print(f"  Verification: {'PASS' if is_valid else 'FAIL'}")
    except Exception as e:
        print(f"  Verification failed: {e}")

    print("\n✅ RFC 3161 integration working")
