"""
Agent Identity - Cryptographically Signed Identity for AI Agents.

Provides:
- Unique agent identity with cryptographic signatures
- Identity verification
- Credential management
- Trust chains
"""
import os
import time
import hmac
import hashlib
import secrets
import base64
from dataclasses import dataclass, field
from typing import Optional, Any
from datetime import datetime, timezone, timedelta


@dataclass
class AgentCredentials:
    """Agent credentials for authentication."""
    agent_id: str
    api_key: str
    secret_key: str
    issued_at: datetime
    expires_at: Optional[datetime] = None
    scopes: list[str] = field(default_factory=list)
    metadata: dict = field(default_factory=dict)
    
    def is_expired(self) -> bool:
        """Check if credentials are expired."""
        if not self.expires_at:
            return False
        return datetime.now(timezone.utc) > self.expires_at
    
    def to_dict(self) -> dict:
        return {
            "agent_id": self.agent_id,
            "api_key": self.api_key,
            "issued_at": self.issued_at.isoformat(),
            "expires_at": self.expires_at.isoformat() if self.expires_at else None,
            "scopes": self.scopes,
        }


@dataclass
class SignedRequest:
    """A cryptographically signed request."""
    request_id: str
    agent_id: str
    timestamp: str
    signature: str
    payload_hash: str
    nonce: str


class AgentIdentity:
    """
    Manages agent identity and cryptographic operations.
    
    Provides:
    - Request signing for authentication
    - Identity verification
    - Credential rotation
    - Trust attestation
    """
    
    SIGNATURE_ALGORITHM = "sha256"
    NONCE_LENGTH = 16
    MAX_TIMESTAMP_DRIFT_SECONDS = 300  # 5 minutes
    
    def __init__(
        self,
        agent_id: str,
        api_key: str,
        secret_key: Optional[str] = None,
        environment: str = "production",
    ):
        """
        Initialize agent identity.
        
        Args:
            agent_id: Unique agent identifier
            api_key: Public API key
            secret_key: Secret key for signing (generated if not provided)
            environment: Deployment environment
        """
        self.agent_id = agent_id
        self.api_key = api_key
        self._secret_key = secret_key or self._generate_secret_key()
        self.environment = environment
        self._request_counter = 0
        self._used_nonces: set[str] = set()
        self._nonce_expiry: dict[str, float] = {}
        
    def _generate_secret_key(self) -> str:
        """Generate a cryptographically secure secret key."""
        return secrets.token_urlsafe(32)
    
    def _generate_nonce(self) -> str:
        """Generate a unique nonce for request signing."""
        return secrets.token_urlsafe(self.NONCE_LENGTH)
    
    def _get_timestamp(self) -> str:
        """Get current timestamp in ISO format."""
        return datetime.now(timezone.utc).isoformat()
    
    def _hash_payload(self, payload: Any) -> str:
        """Create SHA-256 hash of payload."""
        if payload is None:
            payload_str = ""
        elif isinstance(payload, str):
            payload_str = payload
        elif isinstance(payload, bytes):
            payload_str = payload.decode('utf-8')
        else:
            import json
            payload_str = json.dumps(payload, sort_keys=True, default=str)
        
        return hashlib.sha256(payload_str.encode()).hexdigest()
    
    def sign_request(
        self,
        method: str,
        path: str,
        payload: Optional[Any] = None,
        headers: Optional[dict] = None,
    ) -> tuple[dict, SignedRequest]:
        """
        Sign an API request.
        
        Args:
            method: HTTP method
            path: Request path
            payload: Request body
            headers: Additional headers
            
        Returns:
            Tuple of (signed_headers, signature_details)
        """
        timestamp = self._get_timestamp()
        nonce = self._generate_nonce()
        payload_hash = self._hash_payload(payload)
        
        # Build canonical request string
        self._request_counter += 1
        request_id = f"{self.agent_id}-{int(time.time())}-{self._request_counter:06d}"
        
        canonical_string = "\n".join([
            method.upper(),
            path,
            timestamp,
            nonce,
            payload_hash,
            self.agent_id,
        ])
        
        # Create HMAC signature
        signature = hmac.new(
            self._secret_key.encode(),
            canonical_string.encode(),
            hashlib.sha256,
        ).hexdigest()
        
        # Build signed headers
        signed_headers = {
            "X-Vaikora-Agent-ID": self.agent_id,
            "X-Vaikora-Request-ID": request_id,
            "X-Vaikora-Timestamp": timestamp,
            "X-Vaikora-Nonce": nonce,
            "X-Vaikora-Signature": signature,
            "X-Vaikora-Payload-Hash": payload_hash,
        }
        
        if headers:
            signed_headers.update(headers)
        
        signature_details = SignedRequest(
            request_id=request_id,
            agent_id=self.agent_id,
            timestamp=timestamp,
            signature=signature,
            payload_hash=payload_hash,
            nonce=nonce,
        )
        
        # Track nonce to prevent replay
        self._used_nonces.add(nonce)
        self._nonce_expiry[nonce] = time.time() + self.MAX_TIMESTAMP_DRIFT_SECONDS
        self._cleanup_nonces()
        
        return signed_headers, signature_details
    
    def verify_signature(
        self,
        method: str,
        path: str,
        timestamp: str,
        nonce: str,
        payload_hash: str,
        provided_signature: str,
        agent_id: str,
    ) -> tuple[bool, Optional[str]]:
        """
        Verify a signed request.
        
        Args:
            method: HTTP method
            path: Request path
            timestamp: Request timestamp
            nonce: Request nonce
            payload_hash: Hash of payload
            provided_signature: Signature to verify
            agent_id: Agent ID from request
            
        Returns:
            Tuple of (is_valid, error_message)
        """
        # Verify agent ID matches
        if agent_id != self.agent_id:
            return False, "Agent ID mismatch"
        
        # Verify timestamp freshness
        try:
            request_time = datetime.fromisoformat(timestamp.replace("Z", "+00:00"))
            now = datetime.now(timezone.utc)
            drift = abs((now - request_time).total_seconds())
            
            if drift > self.MAX_TIMESTAMP_DRIFT_SECONDS:
                return False, f"Timestamp too old: {drift}s drift"
        except Exception as e:
            return False, f"Invalid timestamp format: {e}"
        
        # Verify nonce hasn't been used (replay prevention)
        if nonce in self._used_nonces:
            return False, "Nonce already used (possible replay attack)"
        
        # Rebuild canonical string and verify signature
        canonical_string = "\n".join([
            method.upper(),
            path,
            timestamp,
            nonce,
            payload_hash,
            agent_id,
        ])
        
        expected_signature = hmac.new(
            self._secret_key.encode(),
            canonical_string.encode(),
            hashlib.sha256,
        ).hexdigest()
        
        if not hmac.compare_digest(expected_signature, provided_signature):
            return False, "Signature mismatch"
        
        # Mark nonce as used
        self._used_nonces.add(nonce)
        self._nonce_expiry[nonce] = time.time() + self.MAX_TIMESTAMP_DRIFT_SECONDS
        
        return True, None
    
    def _cleanup_nonces(self):
        """Clean up expired nonces."""
        now = time.time()
        expired = [n for n, exp in self._nonce_expiry.items() if exp < now]
        for nonce in expired:
            self._used_nonces.discard(nonce)
            del self._nonce_expiry[nonce]
    
    def create_attestation(self, claims: dict) -> str:
        """
        Create a signed attestation of claims.
        
        Args:
            claims: Dictionary of claims to attest
            
        Returns:
            Base64-encoded signed attestation
        """
        import json
        
        attestation = {
            "agent_id": self.agent_id,
            "claims": claims,
            "issued_at": self._get_timestamp(),
            "environment": self.environment,
        }
        
        payload = json.dumps(attestation, sort_keys=True)
        signature = hmac.new(
            self._secret_key.encode(),
            payload.encode(),
            hashlib.sha256,
        ).hexdigest()
        
        attestation["signature"] = signature
        
        return base64.urlsafe_b64encode(
            json.dumps(attestation).encode()
        ).decode()
    
    def verify_attestation(self, encoded_attestation: str) -> tuple[bool, Optional[dict]]:
        """
        Verify a signed attestation.
        
        Args:
            encoded_attestation: Base64-encoded attestation
            
        Returns:
            Tuple of (is_valid, attestation_data or None)
        """
        import json
        
        try:
            attestation = json.loads(
                base64.urlsafe_b64decode(encoded_attestation.encode())
            )
        except Exception:
            return False, None
        
        provided_signature = attestation.pop("signature", None)
        if not provided_signature:
            return False, None
        
        payload = json.dumps(attestation, sort_keys=True)
        expected_signature = hmac.new(
            self._secret_key.encode(),
            payload.encode(),
            hashlib.sha256,
        ).hexdigest()
        
        if not hmac.compare_digest(expected_signature, provided_signature):
            return False, None
        
        return True, attestation
    
    def rotate_secret(self) -> str:
        """
        Rotate the secret key.
        
        Returns:
            New secret key
        """
        old_secret = self._secret_key
        self._secret_key = self._generate_secret_key()
        self._used_nonces.clear()
        self._nonce_expiry.clear()
        return self._secret_key
    
    def get_identity_info(self) -> dict:
        """Get identity information (without secrets)."""
        return {
            "agent_id": self.agent_id,
            "api_key": self.api_key,
            "environment": self.environment,
            "request_count": self._request_counter,
        }
