"""
QWED Attestation Service

Implements the QWED-Attestation specification for cryptographic verification proofs.
Uses JWT with ES256 (ECDSA P-256) for signing attestations.
"""

import hashlib
import json
import base64
import time
import uuid
from dataclasses import dataclass
from typing import Optional, Dict, Any, List, Tuple
from enum import Enum

# Cryptographic imports - using PyJWT with cryptography backend
try:
    import jwt
    from cryptography.hazmat.primitives import serialization
    from cryptography.hazmat.primitives.asymmetric import ec
    from cryptography.hazmat.backends import default_backend
    HAS_CRYPTO = True
except ImportError:
    HAS_CRYPTO = False


class AttestationStatus(Enum):
    """Attestation lifecycle states"""
    ISSUED = "issued"
    VALID = "valid"
    EXPIRED = "expired"
    REVOKED = "revoked"


@dataclass
class VerificationResult:
    """The result of a verification to be attested"""
    status: str  # VERIFIED, FAILED, CORRECTED, BLOCKED
    verified: bool
    engine: str
    confidence: float = 1.0
    query_hash: Optional[str] = None
    proof_hash: Optional[str] = None


@dataclass
class AttestationClaims:
    """QWED Attestation JWT Claims"""
    iss: str  # Issuer DID
    sub: str  # Subject hash
    iat: int  # Issued at
    exp: int  # Expiration
    jti: str  # Attestation ID
    qwed: Dict[str, Any]  # QWED-specific claims


@dataclass
class Attestation:
    """A complete QWED Attestation"""
    jwt_token: str
    claims: AttestationClaims
    header: Dict[str, str]
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            "jwt": self.jwt_token,
            "jti": self.claims.jti,
            "iss": self.claims.iss,
            "iat": self.claims.iat,
            "exp": self.claims.exp,
            "result": self.claims.qwed.get("result", {}),
        }


class IssuerKeyPair:
    """ECDSA P-256 key pair for attestation signing"""
    
    def __init__(self, issuer_did: str, key_id: str):
        if not HAS_CRYPTO:
            raise RuntimeError("cryptography package required for attestations")
        
        self.issuer_did = issuer_did
        self.key_id = key_id
        
        # Generate ECDSA P-256 key pair
        self._private_key = ec.generate_private_key(ec.SECP256R1(), default_backend())
        self._public_key = self._private_key.public_key()
    
    @property
    def private_key_pem(self) -> bytes:
        """Get private key in PEM format"""
        return self._private_key.private_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PrivateFormat.PKCS8,
            encryption_algorithm=serialization.NoEncryption()
        )
    
    @property
    def public_key_pem(self) -> bytes:
        """Get public key in PEM format"""
        return self._public_key.public_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PublicFormat.SubjectPublicKeyInfo
        )
    
    @property
    def jwk(self) -> Dict[str, Any]:
        """Get public key as JWK for verification"""
        from cryptography.hazmat.primitives.asymmetric.utils import decode_dss_signature
        
        # Get the raw numbers from the public key
        numbers = self._public_key.public_numbers()
        
        def int_to_base64url(n: int, length: int) -> str:
            import base64
            return base64.urlsafe_b64encode(
                n.to_bytes(length, 'big')
            ).decode().rstrip('=')
        
        return {
            "kty": "EC",
            "crv": "P-256",
            "x": int_to_base64url(numbers.x, 32),
            "y": int_to_base64url(numbers.y, 32),
            "kid": self.key_id,
        }


class AttestationService:
    """
    Service for creating and verifying QWED attestations.
    
    Implements the QWED-Attestation v1.0 specification.
    """
    

    def __init__(
        self, 
        issuer_did: str = "did:qwed:node:local",
        validity_days: int = 365,
        key_suffix: Optional[str] = None
    ):
        self.issuer_did = issuer_did
        self.validity_days = validity_days
        
        # Key management - deterministic if key_suffix provided
        suffix = key_suffix or "v1"
        self.key_id = f"{issuer_did}#signing-key-{suffix}"
        self._key_pair: Optional[IssuerKeyPair] = None
        
        # Revocation tracking
        self._revoked_attestations: set = set()
        
        # Attestation registry (in-memory, should use DB in production)
        self._attestations: Dict[str, Attestation] = {}
    
    def _ensure_key_pair(self) -> IssuerKeyPair:
        """Lazily initialize key pair"""
        if self._key_pair is None:
            self._key_pair = IssuerKeyPair(self.issuer_did, self.key_id)
        return self._key_pair
    
    def _hash_content(self, content: str) -> str:
        """Create SHA-256 hash of content"""
        return f"sha256:{hashlib.sha256(content.encode()).hexdigest()}"
    
    def create_attestation(
        self,
        verification_result: VerificationResult,
        original_query: str,
        proof_data: Optional[str] = None,
        chain_id: Optional[str] = None,
        chain_index: Optional[int] = None,
    ) -> Attestation:
        """
        Create a signed attestation for a verification result.
        
        Args:
            verification_result: The verification result to attest
            original_query: The original query that was verified
            proof_data: Optional proof data to include
            chain_id: Optional chain ID for linked attestations
            chain_index: Optional index in the chain
        
        Returns:
            Attestation object with signed JWT
        """
        key_pair = self._ensure_key_pair()
        
        now = int(time.time())
        expiry = now + (self.validity_days * 24 * 60 * 60)
        attestation_id = f"att_{uuid.uuid4().hex[:12]}"
        
        # Build QWED-specific claims
        qwed_claims = {
            "version": "1.0",
            "result": {
                "status": verification_result.status,
                "verified": verification_result.verified,
                "engine": verification_result.engine,
                "confidence": verification_result.confidence,
            },
            "query_hash": self._hash_content(original_query),
        }
        
        if proof_data:
            qwed_claims["proof_hash"] = self._hash_content(proof_data)
        
        if chain_id:
            qwed_claims["chain_id"] = chain_id
            if chain_index is not None:
                qwed_claims["chain_index"] = chain_index
        
        # Build full payload
        payload = {
            "iss": self.issuer_did,
            "sub": self._hash_content(original_query),
            "iat": now,
            "exp": expiry,
            "jti": attestation_id,
            "qwed": qwed_claims,
        }
        
        # Build header
        header = {
            "alg": "ES256",
            "typ": "qwed-attestation+jwt",
            "kid": key_pair.key_id,
        }
        
        # Sign the JWT
        token = jwt.encode(
            payload,
            key_pair.private_key_pem,
            algorithm="ES256",
            headers=header,
        )
        
        claims = AttestationClaims(
            iss=payload["iss"],
            sub=payload["sub"],
            iat=payload["iat"],
            exp=payload["exp"],
            jti=attestation_id,
            qwed=qwed_claims,
        )
        
        attestation = Attestation(
            jwt_token=token,
            claims=claims,
            header=header,
        )
        
        # Store in registry
        self._attestations[attestation_id] = attestation
        
        return attestation
    
    def verify_attestation(
        self,
        jwt_token: str,
        trusted_issuers: Optional[List[str]] = None,
    ) -> Tuple[bool, Optional[Dict[str, Any]], Optional[str]]:
        """
        Verify an attestation JWT.
        
        Args:
            jwt_token: The JWT to verify
            trusted_issuers: List of trusted issuer DIDs (None = trust self)
        
        Returns:
            Tuple of (is_valid, claims, error_message)
        """
        if trusted_issuers is None:
            trusted_issuers = [self.issuer_did]
        
        try:
            # Decode without verification first to get issuer
            # Security: We manually decode the payload to get 'iss' and then
            # perform FULL cryptographic verification with the correct key.
            # This avoids using verify_signature=False which triggers security scanners.
            try:
                # Get the payload part (header.payload.signature)
                _, payload_segment, _ = jwt_token.split('.', 2)
                
                # Add padding if needed
                # Correct padding logic: (-n) % 4 gives us the number of '=' needed
                padding = '=' * (-len(payload_segment) % 4)
                payload_data = base64.urlsafe_b64decode(payload_segment + padding)
                unverified = json.loads(payload_data)
                if not isinstance(unverified, dict):
                    raise ValueError("Payload is not a JSON object")
            except Exception:
                return False, None, "Invalid token format"

            issuer = unverified.get("iss")
            
            if issuer not in trusted_issuers:
                safe_issuer = str(issuer)[:128].replace('\n', '').replace('\r', '')
                return False, None, f"Untrusted issuer: {safe_issuer}"
            
            # For self-issued attestations, use our key
            if issuer == self.issuer_did:
                key_pair = self._ensure_key_pair()
                public_key = key_pair.public_key_pem
            else:
                # Would need to resolve DID and get public key
                return False, None, "External issuer key resolution not implemented"
            
            # Verify signature and claims
            claims = jwt.decode(
                jwt_token,
                public_key,
                algorithms=["ES256"],
                options={"require": ["iss", "sub", "iat", "exp", "jti"]},
            )
            
            # Check revocation
            jti = claims.get("jti")
            if jti in self._revoked_attestations:
                return False, None, "Attestation has been revoked"
            
            return True, claims, None
            
        except jwt.ExpiredSignatureError:
            return False, None, "Attestation has expired"
        except jwt.InvalidTokenError as e:
            return False, None, f"Invalid token: {str(e)}"
    
    def revoke_attestation(self, attestation_id: str) -> bool:
        """Revoke an attestation by ID"""
        self._revoked_attestations.add(attestation_id)
        return True
    
    def get_attestation(self, attestation_id: str) -> Optional[Attestation]:
        """Get an attestation by ID"""
        return self._attestations.get(attestation_id)
    
    def get_issuer_info(self) -> Dict[str, Any]:
        """Get issuer information for registry"""
        key_pair = self._ensure_key_pair()
        return {
            "did": self.issuer_did,
            "name": "QWED Local Node",
            "public_keys": [key_pair.jwk],
            "status": "active",
            "certification_level": "basic",
        }


# Singleton instance for the default attestation service
_default_service: Optional[AttestationService] = None


def get_attestation_service() -> AttestationService:
    """Get the default attestation service"""
    global _default_service
    if _default_service is None:
        _default_service = AttestationService()
    return _default_service


def create_verification_attestation(
    status: str,
    verified: bool,
    engine: str,
    query: str,
    confidence: float = 1.0,
    proof_data: Optional[str] = None,
) -> Optional[str]:
    """
    Convenience function to create an attestation for a verification result.
    
    Returns the JWT token string, or None if attestation creation fails.
    """
    try:
        service = get_attestation_service()
        result = VerificationResult(
            status=status,
            verified=verified,
            engine=engine,
            confidence=confidence,
        )
        attestation = service.create_attestation(result, query, proof_data)
        return attestation.jwt_token
    except Exception:
        return None
