#!/usr/bin/env python3

"""

Agent Provenance Chain (APC)

Cryptographic audit trail for autonomous AI agents.

Every action signed. Every decision traceable. Full transparency.

"""

import hashlib

import json

import time

from datetime import datetime

from pathlib import Path

from typing import Any, Dict, Optional

from cryptography.hazmat.primitives.asymmetric.ed25519 import (

    Ed25519PrivateKey,

    Ed25519PublicKey,

)

from cryptography.hazmat.primitives import serialization

import base64

class AgentProvenanceChain:

    """

    Cryptographic audit trail for agent actions.

    Every action is:

    - Timestamped

    - Signed with Ed25519

    - Linked to previous action (blockchain-style)

    - Stored immutably

    """

    def __init__(self, agent_name: str, key_path: Optional[Path] = None):

        self.agent_name = agent_name

        self.key_path = key_path or Path.home() / ".apc" / f"{agent_name}.key"

        self.chain_path = Path.home() / ".apc" / f"{agent_name}.chain.jsonl"

        self._ensure_directories()

        self.private_key = self._load_or_generate_key()

        self.public_key = self.private_key.public_key()

        self.last_hash = self._get_last_hash()

    def _ensure_directories(self):

        """Create necessary directories."""

        self.key_path.parent.mkdir(parents=True, exist_ok=True)

        self.chain_path.parent.mkdir(parents=True, exist_ok=True)

    def _load_or_generate_key(self) -> Ed25519PrivateKey:

        """Load existing key or generate new one."""

        if self.key_path.exists():

            with open(self.key_path, "rb") as f:

                return serialization.load_pem_private_key(f.read(), password=None)

        # Generate new key

        private_key = Ed25519PrivateKey.generate()

        # Save it

        pem = private_key.private_bytes(

            encoding=serialization.Encoding.PEM,

            format=serialization.PrivateFormat.PKCS8,

            encryption_algorithm=serialization.NoEncryption()

        )

        with open(self.key_path, "wb") as f:

            f.write(pem)

        self.key_path.chmod(0o600)

        return private_key

    def _get_last_hash(self) -> str:

        """Get hash of last action in chain."""

        if not self.chain_path.exists():

            return "0" * 64  # Genesis hash

        with open(self.chain_path, "r") as f:

            lines = f.readlines()

            if not lines:

                return "0" * 64

            last_line = lines[-1].strip()

            last_action = json.loads(last_line)

            return last_action["hash"]

    def sign_action(

        self,

        action_type: str,

        payload: Dict[str, Any],

        context: Optional[Dict[str, Any]] = None

    ) -> Dict[str, Any]:

        """

        Sign an action and append to chain.

        Args:

            action_type: Type of action (e.g., "exec", "write_file", "api_call")

            payload: Action data

            context: Optional context (reasoning, session_id, etc.)

        Returns:

            Signed action record

        """

        timestamp = time.time()

        # Build action record

        action = {

            "agent": self.agent_name,

            "timestamp": timestamp,

            "iso_time": datetime.utcfromtimestamp(timestamp).isoformat() + "Z",

            "type": action_type,

            "payload": payload,

            "context": context or {},

            "previous_hash": self.last_hash,

        }

        # Compute hash

        action_bytes = json.dumps(action, sort_keys=True).encode()

        action_hash = hashlib.sha256(action_bytes).hexdigest()

        action["hash"] = action_hash

        # Sign the hash

        signature = self.private_key.sign(action_hash.encode())

        action["signature"] = base64.b64encode(signature).decode()

        # Append to chain

        with open(self.chain_path, "a") as f:

            f.write(json.dumps(action) + "\n")

        # Update last hash

        self.last_hash = action_hash

        return action

    def get_public_key_pem(self) -> str:

        """Get public key in PEM format for verification."""

        pem = self.public_key.public_bytes(

            encoding=serialization.Encoding.PEM,

            format=serialization.PublicFormat.SubjectPublicKeyInfo

        )

        return pem.decode()

    def get_chain(self, limit: Optional[int] = None) -> list:

        """Retrieve action chain."""

        if not self.chain_path.exists():

            return []

        with open(self.chain_path, "r") as f:

            lines = f.readlines()

        if limit:

            lines = lines[-limit:]

        return [json.loads(line) for line in lines if line.strip()]

    def verify_chain_integrity(self) -> tuple[bool, Optional[str]]:

        """

        Verify the entire chain is intact and unmodified.

        Returns:

            (is_valid, error_message)

        """

        chain = self.get_chain()

        if not chain:

            return True, None

        expected_prev = "0" * 64

        for i, action in enumerate(chain):

            # Check previous hash linkage

            if action["previous_hash"] != expected_prev:

                return False, f"Chain broken at index {i}: hash mismatch"

            # Verify signature

            try:

                sig_bytes = base64.b64decode(action["signature"])

                self.public_key.verify(sig_bytes, action["hash"].encode())

            except Exception as e:

                return False, f"Invalid signature at index {i}: {e}"

            expected_prev = action["hash"]

        return True, None

# Convenience function

def create_agent_chain(agent_name: str) -> AgentProvenanceChain:

    """Create or load an agent provenance chain."""

    return AgentProvenanceChain(agent_name)