#!/usr/bin/env python3
"""
receipt_signing.py - Signed Receipts + Determinism Seal

WO-TRUTHY-METRICS-01: Receipt signing utilities for audit-grade receipts.

Integrates with existing tools/sign_receipts.py key management.

Features:
- Ed25519 signatures for payload authenticity
- Determinism seals (environment + seed block digests)
- Structured failure diagnostics

Usage:
    from helix_cdc.utils.receipt_signing import (
        SignedReceipt, sign_receipt,
        compute_environment_digest, compute_seed_block_digest
    )

    # Create receipt
    receipt = SignedReceipt(
        schema="helix:benchmark:v1",
        payload={"result": 42},
        environment_digest=compute_environment_digest(),
        seed_block_digest=compute_seed_block_digest({"input": "test"}),
    )

    # Sign it (uses ~/.helix_keys/signing_key.pem if available)
    signed = sign_receipt(receipt)
    print(json.dumps(signed.to_dict(), indent=2))
"""

from __future__ import annotations

import hashlib
import json
import os
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, List, Optional


@dataclass
class FailureDiagnostic:
    """
    Structured failure info for debugging.

    WO-TRUTHY-METRICS-01: Standardized failure diagnostics for receipts.
    """
    failure_type: str              # "hash_mismatch", "timeout", "exception", etc.
    first_diff_offset: Optional[int] = None
    expected_value: Optional[str] = None
    actual_value: Optional[str] = None
    context: Optional[Dict[str, Any]] = None

    def to_dict(self) -> Dict[str, Any]:
        result = {"failure_type": self.failure_type}

        if self.first_diff_offset is not None:
            result["first_diff"] = {
                "offset": self.first_diff_offset,
                "expected": self.expected_value,
                "actual": self.actual_value,
            }

        if self.context:
            result["context"] = self.context

        return result


@dataclass
class SignedReceipt:
    """
    Wrapper that adds signature + seal to any receipt dict.

    WO-TRUTHY-METRICS-01: Audit-grade receipts with determinism seals.
    """
    schema: str
    payload: Dict[str, Any]

    # Determinism seal
    environment_digest: str = ""   # SHA256 of sorted env vars that affect output
    seed_block_digest: str = ""    # SHA256 of all inputs that determine output

    # Signature fields (filled by sign_receipt)
    signature_b64: Optional[str] = None
    sig_alg: str = "ed25519"
    kid: Optional[str] = None      # Key ID (first 12 chars of SHA256(pubkey))
    signed_at: Optional[str] = None

    # Optional failure diagnostic
    failure: Optional[FailureDiagnostic] = None

    def compute_payload_hash(self) -> str:
        """Deterministic hash of payload for signing."""
        canonical = json.dumps(self.payload, sort_keys=True, separators=(',', ':'))
        return hashlib.sha256(canonical.encode()).hexdigest()

    def to_dict(self) -> Dict[str, Any]:
        result = {
            "schema": self.schema,
            "payload": self.payload,
            "seal": {
                "environment_digest": self.environment_digest,
                "seed_block_digest": self.seed_block_digest,
                "payload_sha256": self.compute_payload_hash(),
            },
        }

        if self.signature_b64:
            result["signature"] = {
                "algorithm": self.sig_alg,
                "kid": self.kid,
                "value": self.signature_b64,
                "signed_at": self.signed_at,
            }
        else:
            result["signature"] = None

        if self.failure:
            result["failure"] = self.failure.to_dict()

        return result


def compute_environment_digest() -> str:
    """
    Hash environment variables that affect determinism.

    Captures HELIX_, CUDA_, OMP_, MKL_ prefixed variables.
    Returns first 16 chars of SHA256 for compactness.
    """
    relevant_vars = sorted([
        f"{k}={v}" for k, v in os.environ.items()
        if k.startswith(("HELIX_", "CUDA_", "OMP_", "MKL_", "TORCH_"))
    ])
    return hashlib.sha256("\n".join(relevant_vars).encode()).hexdigest()[:16]


def compute_seed_block_digest(inputs: Dict[str, Any]) -> str:
    """
    Hash all inputs that determine output.

    Args:
        inputs: Dict of all inputs that affect the output (file paths, params, etc.)

    Returns:
        Full SHA256 hex digest of canonical JSON representation.
    """
    canonical = json.dumps(inputs, sort_keys=True, separators=(',', ':'))
    return hashlib.sha256(canonical.encode()).hexdigest()


def sign_receipt(receipt: SignedReceipt, key_dir: Optional[Path] = None) -> SignedReceipt:
    """
    Sign a receipt using Ed25519 key from .helix_keys/.

    Looks for keys in:
    1. key_dir parameter (if provided)
    2. ~/.helix_keys/
    3. .helix_keys/ in current directory

    If no key is found, returns the receipt unsigned.

    Args:
        receipt: SignedReceipt to sign
        key_dir: Optional path to directory containing signing_key.pem

    Returns:
        SignedReceipt with signature fields filled (if key available)
    """
    # Search for key in multiple locations
    search_paths = []
    if key_dir:
        search_paths.append(key_dir)
    search_paths.extend([
        Path.home() / ".helix_keys",
        Path(".helix_keys"),
        Path("keys"),  # tools/sign_receipts.py default
    ])

    key_path = None
    for search_dir in search_paths:
        candidate = search_dir / "ed25519_priv.pem"
        if candidate.exists():
            key_path = candidate
            break
        # Also check for signing_key.pem (alternative name)
        candidate = search_dir / "signing_key.pem"
        if candidate.exists():
            key_path = candidate
            break

    if not key_path:
        # No key available - return unsigned
        return receipt

    try:
        from cryptography.hazmat.primitives import serialization
        from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey
        import base64

        with open(key_path, "rb") as f:
            private_key = serialization.load_pem_private_key(f.read(), password=None)

        # Compute what we're signing: schema + payload hash
        payload_hash = receipt.compute_payload_hash()
        message = f"{receipt.schema}:{payload_hash}".encode()

        # Sign
        signature = private_key.sign(message)

        # Compute key ID from public key
        public_key_bytes = private_key.public_key().public_bytes(
            encoding=serialization.Encoding.Raw,
            format=serialization.PublicFormat.Raw
        )
        kid = hashlib.sha256(public_key_bytes).hexdigest()[:12]

        # Update receipt
        receipt.signature_b64 = base64.b64encode(signature).decode()
        receipt.kid = kid
        receipt.signed_at = datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")

    except ImportError:
        # cryptography not installed
        pass
    except Exception as e:
        # Signing failed - log but don't crash
        print(f"[WARN] Receipt signing failed: {e}")

    return receipt


def verify_receipt(receipt: SignedReceipt, pubkey_path: Optional[Path] = None) -> bool:
    """
    Verify a signed receipt.

    Args:
        receipt: SignedReceipt with signature to verify
        pubkey_path: Path to public key PEM file

    Returns:
        True if signature is valid, False otherwise
    """
    if not receipt.signature_b64:
        return False

    # Search for public key
    search_paths = []
    if pubkey_path:
        search_paths.append(pubkey_path)
    search_paths.extend([
        Path.home() / ".helix_keys" / "ed25519_pub.pem",
        Path(".helix_keys") / "ed25519_pub.pem",
        Path("keys") / "ed25519_pub.pem",
    ])

    key_path = None
    for candidate in search_paths:
        if candidate.exists():
            key_path = candidate
            break

    if not key_path:
        return False

    try:
        from cryptography.hazmat.primitives import serialization
        import base64

        with open(key_path, "rb") as f:
            public_key = serialization.load_pem_public_key(f.read())

        # Recompute message
        payload_hash = receipt.compute_payload_hash()
        message = f"{receipt.schema}:{payload_hash}".encode()

        # Verify
        signature = base64.b64decode(receipt.signature_b64)
        public_key.verify(signature, message)
        return True

    except Exception:
        return False


def create_failure_diagnostic(
    failure_type: str,
    expected: Optional[bytes] = None,
    actual: Optional[bytes] = None,
    context: Optional[Dict[str, Any]] = None,
) -> FailureDiagnostic:
    """
    Create a failure diagnostic with first-diff detection.

    Args:
        failure_type: Type of failure (hash_mismatch, timeout, exception)
        expected: Expected bytes (for hash_mismatch)
        actual: Actual bytes received
        context: Additional context dict

    Returns:
        FailureDiagnostic with first_diff populated if applicable
    """
    diagnostic = FailureDiagnostic(
        failure_type=failure_type,
        context=context,
    )

    if expected is not None and actual is not None and failure_type == "hash_mismatch":
        # Find first difference
        for i, (e, a) in enumerate(zip(expected, actual)):
            if e != a:
                diagnostic.first_diff_offset = i
                diagnostic.expected_value = hex(e)
                diagnostic.actual_value = hex(a)
                break
        else:
            # Lengths differ
            if len(expected) != len(actual):
                diagnostic.first_diff_offset = min(len(expected), len(actual))
                diagnostic.expected_value = f"len={len(expected)}"
                diagnostic.actual_value = f"len={len(actual)}"

    return diagnostic


if __name__ == "__main__":
    # Self-test
    print("Receipt Signing Utility - Self Test")
    print("=" * 50)

    # Create test receipt
    receipt = SignedReceipt(
        schema="helix:test:v1",
        payload={"test_key": "test_value", "number": 42},
        environment_digest=compute_environment_digest(),
        seed_block_digest=compute_seed_block_digest({"input": "test_input"}),
    )

    print("Created unsigned receipt:")
    print(f"  Schema: {receipt.schema}")
    print(f"  Payload hash: {receipt.compute_payload_hash()[:16]}...")
    print(f"  Env digest: {receipt.environment_digest}")
    print(f"  Seed digest: {receipt.seed_block_digest[:16]}...")

    # Try to sign
    signed = sign_receipt(receipt)

    if signed.signature_b64:
        print(f"\nSigned successfully:")
        print(f"  Key ID: {signed.kid}")
        print(f"  Signed at: {signed.signed_at}")
        print(f"  Signature: {signed.signature_b64[:32]}...")

        # Verify
        if verify_receipt(signed):
            print("  Verification: PASS")
        else:
            print("  Verification: FAIL")
    else:
        print("\nNo signing key found - receipt is unsigned")
        print("Generate keys with: python tools/sign_receipts.py gen")

    print("\nFull receipt dict:")
    import json
    print(json.dumps(signed.to_dict(), indent=2))
