"""
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 time
import uuid
from dataclasses import dataclass, field
from datetime import datetime, timedelta
from typing import Optional, Dict, Any, List
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
    ):
        self.issuer_did = issuer_did
        self.validity_days = validity_days
        
        # Key management
        self.key_id = f"{issuer_did}#signing-key-{datetime.now().year}"
        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
            unverified = jwt.decode(jwt_token, options={"verify_signature": False})
            issuer = unverified.get("iss")
            
            if issuer not in trusted_issuers:
                return False, None, f"Untrusted issuer: {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
