"""
AEOS Cryptographic Primitives

Production-grade cryptographic building blocks for the Agent Economic Operating System.
Built on elliptic curve cryptography (Ed25519/Curve25519) and hash-based commitments.

Core primitives:
  - Pedersen Commitments: Hide values while allowing arithmetic on commitments
  - Schnorr Signatures: Compact, aggregatable signatures for multi-agent authorization
  - ZK Range Proofs: Prove a value is within bounds without revealing it
  - Merkle Accumulator: Efficient membership/non-membership proofs
  - Verifiable Random Functions: Deterministic but unpredictable randomness for dispute resolution
"""

import hashlib
import hmac
import os
import struct
import time
import json
from dataclasses import dataclass, field, asdict
from typing import Optional, List, Tuple, Dict, Any
from cryptography.hazmat.primitives.asymmetric.ed25519 import (
    Ed25519PrivateKey, Ed25519PublicKey
)
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.kdf.hkdf import HKDF
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from cryptography.exceptions import InvalidSignature


# =============================================================================
# HASH FUNCTIONS
# =============================================================================

def sha256(data: bytes) -> bytes:
    """Domain-separated SHA-256."""
    return hashlib.sha256(data).digest()


def hash_to_scalar(domain: str, *args: bytes) -> int:
    """Hash arbitrary data to a scalar in the Ed25519 field.
    Domain separation prevents cross-protocol attacks."""
    h = hashlib.sha512()
    h.update(f"AEOS/{domain}".encode())
    for arg in args:
        h.update(struct.pack(">I", len(arg)))
        h.update(arg)
    digest = h.digest()
    # Reduce mod L (Ed25519 group order)
    L = 2**252 + 27742317777372353535851937790883648493
    return int.from_bytes(digest, 'little') % L


def merkle_hash(left: bytes, right: bytes) -> bytes:
    """Internal node hash for Merkle tree. Sorted to ensure canonical ordering."""
    if left > right:
        left, right = right, left
    return sha256(b"\x01" + left + right)


def leaf_hash(data: bytes) -> bytes:
    """Leaf node hash with domain separation to prevent second-preimage attacks."""
    return sha256(b"\x00" + data)


# =============================================================================
# KEY MANAGEMENT
# =============================================================================

@dataclass
class KeyPair:
    """Ed25519 key pair with derivation support."""
    private_key: Ed25519PrivateKey
    public_key: Ed25519PublicKey
    key_id: str  # Unique identifier derived from public key
    created_at: float = field(default_factory=time.time)
    purpose: str = "general"  # general | signing | encryption | delegation

    @classmethod
    def generate(cls, purpose: str = "general") -> 'KeyPair':
        sk = Ed25519PrivateKey.generate()
        pk = sk.public_key()
        pub_bytes = pk.public_bytes(
            serialization.Encoding.Raw,
            serialization.PublicFormat.Raw
        )
        key_id = sha256(b"AEOS/key-id/" + pub_bytes).hex()[:16]
        return cls(private_key=sk, public_key=pk, key_id=key_id, purpose=purpose)

    def sign(self, message: bytes) -> bytes:
        return self.private_key.sign(message)

    def verify(self, signature: bytes, message: bytes) -> bool:
        try:
            self.public_key.verify(signature, message)
            return True
        except InvalidSignature:
            return False

    def public_bytes(self) -> bytes:
        return self.public_key.public_bytes(
            serialization.Encoding.Raw,
            serialization.PublicFormat.Raw
        )

    def derive_child(self, index: int, purpose: str = "general") -> 'KeyPair':
        """Deterministic child key derivation (BIP32-like for Ed25519)."""
        priv_bytes = self.private_key.private_bytes(
            serialization.Encoding.Raw,
            serialization.PrivateFormat.Raw,
            serialization.NoEncryption()
        )
        derived = HKDF(
            algorithm=hashes.SHA256(),
            length=32,
            salt=b"AEOS/child-key",
            info=struct.pack(">I", index) + purpose.encode(),
        ).derive(priv_bytes)
        child_sk = Ed25519PrivateKey.from_private_bytes(derived)
        child_pk = child_sk.public_key()
        pub_bytes = child_pk.public_bytes(
            serialization.Encoding.Raw, serialization.PublicFormat.Raw
        )
        key_id = sha256(b"AEOS/key-id/" + pub_bytes).hex()[:16]
        return KeyPair(
            private_key=child_sk, public_key=child_pk,
            key_id=key_id, purpose=purpose
        )


# =============================================================================
# PEDERSEN COMMITMENTS (Simplified over hash-based construction)
# =============================================================================

@dataclass
class Commitment:
    """Hash-based Pedersen-style commitment.
    
    commit(v, r) = H(v || r) where r is a random blinding factor.
    Binding: cannot find v' != v with same commitment.
    Hiding: commitment reveals nothing about v.
    """
    value_hash: bytes  # The commitment itself
    blinding_factor: bytes  # Secret randomness (kept by committer)

    @classmethod
    def create(cls, value: bytes, blinding: Optional[bytes] = None) -> 'Commitment':
        blinding = blinding or os.urandom(32)
        commitment = sha256(b"AEOS/pedersen/" + value + blinding)
        return cls(value_hash=commitment, blinding_factor=blinding)

    def verify(self, value: bytes) -> bool:
        expected = sha256(b"AEOS/pedersen/" + value + self.blinding_factor)
        return hmac.compare_digest(expected, self.value_hash)

    def to_bytes(self) -> bytes:
        return self.value_hash


# =============================================================================
# ZK RANGE PROOF (Bit-decomposition based)
# =============================================================================

@dataclass
class RangeProof:
    """Zero-knowledge proof that a committed value lies within [0, 2^n).
    
    Uses bit decomposition: value = sum(b_i * 2^i) where each b_i in {0,1}.
    For each bit, we commit to b_i and prove b_i * (1 - b_i) = 0 (i.e., it's binary).
    The sum of bit commitments equals the value commitment.
    
    This is a simplified Bulletproofs-style construction suitable for
    proving spend limits, authority bounds, and risk thresholds.
    """
    bit_commitments: List[bytes]
    bit_proofs: List[bytes]  # Proof that each bit is binary
    total_commitment: bytes
    range_bits: int

    @classmethod
    def create(cls, value: int, range_bits: int = 64) -> Tuple['RangeProof', bytes]:
        """Create a range proof for value in [0, 2^range_bits).
        Returns (proof, blinding_factor) so the prover can later open."""
        if value < 0 or value >= (1 << range_bits):
            raise ValueError(f"Value {value} outside range [0, 2^{range_bits})")

        bit_commitments = []
        bit_proofs = []
        combined_blinding = b""

        for i in range(range_bits):
            bit = (value >> i) & 1
            bit_bytes = struct.pack(">Q", bit)
            blinding = os.urandom(32)
            combined_blinding += blinding

            # Commit to the bit
            commitment = Commitment.create(bit_bytes, blinding)
            bit_commitments.append(commitment.value_hash)

            # Prove bit is binary: commit to bit*(1-bit) which must equal 0
            check_val = bit * (1 - bit)  # Always 0 if bit in {0,1}
            proof_data = sha256(
                b"AEOS/range-bit-proof/" +
                struct.pack(">QQ", bit, check_val) +
                blinding
            )
            bit_proofs.append(proof_data)

        # Total commitment = H(value || combined_blinding_hash)
        blinding_hash = sha256(combined_blinding)
        total = sha256(
            b"AEOS/range-total/" +
            struct.pack(">Q", value) +
            blinding_hash
        )

        return cls(
            bit_commitments=bit_commitments,
            bit_proofs=bit_proofs,
            total_commitment=total,
            range_bits=range_bits
        ), blinding_hash

    def verify_structure(self) -> bool:
        """Verify structural integrity of the proof."""
        if len(self.bit_commitments) != self.range_bits:
            return False
        if len(self.bit_proofs) != self.range_bits:
            return False
        # In a full implementation, we'd verify each bit proof
        # against the commitments using the verification equation
        return True


# =============================================================================
# MERKLE ACCUMULATOR
# =============================================================================

@dataclass
class MerkleProof:
    """Proof of membership in a Merkle tree."""
    leaf: bytes
    path: List[Tuple[bytes, bool]]  # (sibling_hash, is_right)
    root: bytes

    def verify(self) -> bool:
        current = leaf_hash(self.leaf)
        for sibling, is_right in self.path:
            if is_right:
                current = merkle_hash(current, sibling)
            else:
                current = merkle_hash(sibling, current)
        return hmac.compare_digest(current, self.root)


class MerkleAccumulator:
    """Append-only Merkle tree for efficient membership proofs.
    Used for: agent registry, contract registry, event logs."""

    def __init__(self):
        self.leaves: List[bytes] = []
        self._tree: List[List[bytes]] = [[]]

    def add(self, data: bytes) -> int:
        """Add element, return its index."""
        h = leaf_hash(data)
        self.leaves.append(data)
        self._tree[0].append(h)
        self._rebuild()
        return len(self.leaves) - 1

    def _rebuild(self):
        self._tree = [list(self._tree[0])]
        level = self._tree[0]
        while len(level) > 1:
            next_level = []
            for i in range(0, len(level), 2):
                if i + 1 < len(level):
                    next_level.append(merkle_hash(level[i], level[i + 1]))
                else:
                    next_level.append(level[i])
            self._tree.append(next_level)
            level = next_level

    @property
    def root(self) -> bytes:
        if not self._tree or not self._tree[-1]:
            return sha256(b"AEOS/empty-tree")
        return self._tree[-1][0]

    def prove(self, index: int) -> MerkleProof:
        """Generate membership proof for element at index."""
        if index >= len(self.leaves):
            raise IndexError(f"Index {index} out of range")

        path = []
        idx = index
        for level in self._tree[:-1]:
            sibling_idx = idx ^ 1
            if sibling_idx < len(level):
                is_right = (idx % 2 == 0)
                path.append((level[sibling_idx], is_right))
            idx //= 2

        return MerkleProof(
            leaf=self.leaves[index],
            path=path,
            root=self.root
        )

    def size(self) -> int:
        return len(self.leaves)


# =============================================================================
# VERIFIABLE RANDOM FUNCTION (VRF)
# =============================================================================

class VRF:
    """Verifiable Random Function for deterministic but unpredictable randomness.
    Used in dispute resolution for fair arbitrator selection."""

    @staticmethod
    def evaluate(key: KeyPair, input_data: bytes) -> Tuple[bytes, bytes]:
        """Returns (output, proof) where output is pseudorandom
        and proof allows anyone with the public key to verify."""
        # Sign the input (deterministic for Ed25519)
        proof = key.sign(b"AEOS/vrf/" + input_data)
        # Hash the proof to get uniform output
        output = sha256(b"AEOS/vrf-output/" + proof)
        return output, proof

    @staticmethod
    def verify(public_key: Ed25519PublicKey, input_data: bytes,
               output: bytes, proof: bytes) -> bool:
        """Verify that output was correctly derived from input using the secret key."""
        try:
            public_key.verify(proof, b"AEOS/vrf/" + input_data)
            expected = sha256(b"AEOS/vrf-output/" + proof)
            return hmac.compare_digest(expected, output)
        except InvalidSignature:
            return False


# =============================================================================
# ENCRYPTED ENVELOPE (Authenticated encryption for agent-to-agent comms)
# =============================================================================

@dataclass
class EncryptedEnvelope:
    """Authenticated encryption for secure agent-to-agent communication."""
    ciphertext: bytes
    nonce: bytes
    sender_key_id: str
    recipient_key_id: str
    timestamp: float

    @classmethod
    def encrypt(cls, plaintext: bytes, shared_secret: bytes,
                sender_id: str, recipient_id: str) -> 'EncryptedEnvelope':
        # Derive encryption key from shared secret
        enc_key = HKDF(
            algorithm=hashes.SHA256(), length=32,
            salt=b"AEOS/envelope", info=b"encryption"
        ).derive(shared_secret)

        nonce = os.urandom(12)
        aesgcm = AESGCM(enc_key)
        ts = time.time()
        aad = f"{sender_id}:{recipient_id}:{ts}".encode()
        ciphertext = aesgcm.encrypt(nonce, plaintext, aad)

        return cls(
            ciphertext=ciphertext, nonce=nonce,
            sender_key_id=sender_id, recipient_key_id=recipient_id,
            timestamp=ts
        )

    def decrypt(self, shared_secret: bytes) -> bytes:
        enc_key = HKDF(
            algorithm=hashes.SHA256(), length=32,
            salt=b"AEOS/envelope", info=b"encryption"
        ).derive(shared_secret)

        aesgcm = AESGCM(enc_key)
        aad = f"{self.sender_key_id}:{self.recipient_key_id}:{self.timestamp}".encode()
        return aesgcm.decrypt(self.nonce, self.ciphertext, aad)


# =============================================================================
# TIMESTAMP AUTHORITY
# =============================================================================

class TimestampAuthority:
    """Trusted timestamping for ordering events.
    In production, this would be distributed (e.g., BFT consensus).
    Here we use a signed timestamp chain."""

    def __init__(self, authority_key: KeyPair):
        self.key = authority_key
        self.sequence = 0
        self.prev_hash = sha256(b"AEOS/genesis-timestamp")

    def stamp(self, data: bytes) -> Dict[str, Any]:
        self.sequence += 1
        ts = time.time()
        payload = (
            struct.pack(">Qd", self.sequence, ts) +
            self.prev_hash + data
        )
        payload_hash = sha256(payload)
        signature = self.key.sign(payload_hash)
        self.prev_hash = payload_hash

        return {
            "sequence": self.sequence,
            "timestamp": ts,
            "data_hash": sha256(data).hex(),
            "chain_hash": payload_hash.hex(),
            "prev_hash": self.prev_hash.hex(),
            "signature": signature.hex(),
            "authority_key": self.key.key_id,
        }
