"""
Identity Wallet for Anchor8 SDK

Provides client-side key management and request signing for
Universal Agent Identity (UAI) integration.

Features:
- Ed25519 key pair generation
- Request signing for DID-Auth
- Credential storage and presentation
"""

import os
import json
import base64
import logging
from typing import Dict, Any, Optional, Tuple
from pathlib import Path

from cryptography.hazmat.primitives.asymmetric.ed25519 import (
    Ed25519PrivateKey,
    Ed25519PublicKey
)
from cryptography.hazmat.primitives import serialization

logger = logging.getLogger("anchor8.sdk.identity")


class IdentityWallet:
    """
    Client-side identity wallet for managing agent keys and credentials.
    
    Handles:
    - Key pair generation and storage
    - Request signing for DID-Auth
    - Credential storage and presentation
    """
    
    def __init__(
        self,
        private_key_jwk: Optional[Dict[str, str]] = None,
        did: Optional[str] = None,
        key_file: Optional[str] = None
    ):
        """
        Initialize the identity wallet.
        
        Args:
            private_key_jwk: Existing private key in JWK format
            did: The agent's DID
            key_file: Path to file containing the private key
        """
        self.did = did
        self._private_key: Optional[Ed25519PrivateKey] = None
        self._public_key: Optional[Ed25519PublicKey] = None
        self._credentials: Dict[str, Dict] = {}  # credential_id -> credential
        
        # Load or generate keys
        if private_key_jwk:
            self._load_from_jwk(private_key_jwk)
        elif key_file:
            self._load_from_file(key_file)
        elif os.getenv("ANCHOR8_AGENT_KEY"):
            # Load from environment variable (JSON-encoded JWK)
            jwk = json.loads(os.getenv("ANCHOR8_AGENT_KEY"))
            self._load_from_jwk(jwk)
        else:
            # Generate new key pair
            self._generate_keypair()
    
    def _generate_keypair(self) -> None:
        """Generate a new Ed25519 key pair."""
        self._private_key = Ed25519PrivateKey.generate()
        self._public_key = self._private_key.public_key()
        logger.info("Generated new Ed25519 key pair")
    
    def _load_from_jwk(self, jwk: Dict[str, str]) -> None:
        """Load private key from JWK format."""
        d_bytes = base64.urlsafe_b64decode(jwk["d"] + "==")
        self._private_key = Ed25519PrivateKey.from_private_bytes(d_bytes)
        self._public_key = self._private_key.public_key()
        logger.info("Loaded key from JWK")
    
    def _load_from_file(self, key_file: str) -> None:
        """Load private key from a JSON file."""
        path = Path(key_file)
        if not path.exists():
            logger.warning(f"Key file not found: {key_file}, generating new key")
            self._generate_keypair()
            self.save_to_file(key_file)
            return
        
        with open(path) as f:
            data = json.load(f)
        
        if "private_key" in data:
            self._load_from_jwk(data["private_key"])
        if "did" in data:
            self.did = data["did"]
        if "credentials" in data:
            self._credentials = data["credentials"]
    
    def save_to_file(self, key_file: str) -> None:
        """Save wallet to a JSON file."""
        path = Path(key_file)
        path.parent.mkdir(parents=True, exist_ok=True)
        
        data = {
            "did": self.did,
            "private_key": self.get_private_key_jwk(),
            "public_key": self.get_public_key_jwk(),
            "credentials": self._credentials
        }
        
        with open(path, "w") as f:
            json.dump(data, f, indent=2)
        
        logger.info(f"Saved wallet to {key_file}")
    
    def get_private_key_jwk(self) -> Dict[str, str]:
        """Get private key in JWK format."""
        if not self._private_key:
            raise ValueError("No private key available")
        
        private_bytes = self._private_key.private_bytes(
            encoding=serialization.Encoding.Raw,
            format=serialization.PrivateFormat.Raw,
            encryption_algorithm=serialization.NoEncryption()
        )
        public_bytes = self._public_key.public_bytes(
            encoding=serialization.Encoding.Raw,
            format=serialization.PublicFormat.Raw
        )
        
        return {
            "kty": "OKP",
            "crv": "Ed25519",
            "x": base64.urlsafe_b64encode(public_bytes).rstrip(b"=").decode(),
            "d": base64.urlsafe_b64encode(private_bytes).rstrip(b"=").decode()
        }
    
    def get_public_key_jwk(self) -> Dict[str, str]:
        """Get public key in JWK format."""
        if not self._public_key:
            raise ValueError("No public key available")
        
        public_bytes = self._public_key.public_bytes(
            encoding=serialization.Encoding.Raw,
            format=serialization.PublicFormat.Raw
        )
        
        return {
            "kty": "OKP",
            "crv": "Ed25519",
            "x": base64.urlsafe_b64encode(public_bytes).rstrip(b"=").decode()
        }
    
    def sign(self, data: bytes) -> str:
        """
        Sign data with the private key.
        
        Args:
            data: Bytes to sign
            
        Returns:
            Base64-encoded signature
        """
        if not self._private_key:
            raise ValueError("No private key available for signing")
        
        signature = self._private_key.sign(data)
        return base64.urlsafe_b64encode(signature).rstrip(b"=").decode()
    
    def sign_request(self, payload: Dict[str, Any]) -> Tuple[Dict[str, str], bytes]:
        """
        Sign a request payload for DID-Auth.
        
        Args:
            payload: The request payload to sign
            
        Returns:
            Tuple of (headers, body) where headers contain X-Agent-DID and X-Signature
        """
        body = json.dumps(payload, sort_keys=True).encode()
        signature = self.sign(body)
        
        headers = {
            "X-Agent-DID": self.did or "did:key:unsigned",
            "X-Signature": signature
        }
        
        return headers, body
    
    def store_credential(self, credential: Dict[str, Any]) -> None:
        """
        Store a credential in the wallet.
        
        Args:
            credential: The Verifiable Credential to store
        """
        credential_id = credential.get("id")
        if not credential_id:
            raise ValueError("Credential must have an 'id' field")
        
        self._credentials[credential_id] = credential
        logger.debug(f"Stored credential: {credential_id}")
    
    def get_credential(self, credential_id: str) -> Optional[Dict[str, Any]]:
        """Get a stored credential by ID."""
        return self._credentials.get(credential_id)
    
    def list_credentials(self) -> list:
        """List all stored credential IDs."""
        return list(self._credentials.keys())
    
    def create_presentation(
        self,
        credential_ids: list,
        challenge: Optional[str] = None
    ) -> Dict[str, Any]:
        """
        Create a Verifiable Presentation from stored credentials.
        
        Args:
            credential_ids: List of credential IDs to include
            challenge: Optional nonce for replay protection
            
        Returns:
            Signed Verifiable Presentation
        """
        credentials = []
        for cred_id in credential_ids:
            cred = self._credentials.get(cred_id)
            if cred:
                credentials.append(cred)
        
        presentation = {
            "@context": ["https://www.w3.org/2018/credentials/v1"],
            "type": ["VerifiablePresentation"],
            "holder": self.did,
            "verifiableCredential": credentials
        }
        
        if challenge:
            presentation["challenge"] = challenge
        
        # Sign the presentation
        presentation_bytes = json.dumps(presentation, sort_keys=True).encode()
        signature = self.sign(presentation_bytes)
        
        presentation["proof"] = {
            "type": "Ed25519Signature2020",
            "verificationMethod": f"{self.did}#key-1",
            "proofPurpose": "authentication",
            "proofValue": signature
        }
        
        return presentation
    
    @property
    def has_identity(self) -> bool:
        """Check if the wallet has a configured identity."""
        return self._private_key is not None and self.did is not None
