"""SD-JWT attestations for agent trust and reputation.

Agents can issue attestations about other agents (e.g. "completed task
with score 0.95") using Selective Disclosure JWTs (SD-JWT, RFC 9901).
Verifiers can check claims without seeing all details.

Usage::

    from skytale_sdk.attestation import Attestation, create_attestation, verify_attestation

    # Issuer creates an attestation
    att = create_attestation(
        issuer_identity=my_identity,
        subject_did="did:key:z6Mk...",
        claims={"task_score": 0.95, "task_type": "analysis"},
        disclosed=["task_score"],  # Only this claim is visible to verifiers
    )

    # Verifier checks it
    result = verify_attestation(att, issuer_identity.public_key)
"""

from __future__ import annotations

import hashlib
import json
import os
import time
from dataclasses import dataclass, field
from typing import Optional


@dataclass
class Attestation:
    """An SD-JWT attestation about an agent.

    Attributes:
        issuer: DID of the issuing agent.
        subject: DID of the subject agent.
        claims: All claims (some may be selectively disclosed).
        disclosed_keys: Which claim keys are disclosed to verifiers.
        issued_at: Unix timestamp.
        expires_at: Optional expiry timestamp.
        signature: Ed25519 signature over the attestation payload.
        salt_map: Map of claim keys to their random salts (for SD-JWT).
    """

    issuer: str
    subject: str
    claims: dict
    disclosed_keys: list[str] = field(default_factory=list)
    issued_at: float = field(default_factory=time.time)
    expires_at: Optional[float] = None
    signature: Optional[bytes] = None
    salt_map: dict[str, str] = field(default_factory=dict)


def create_attestation(
    issuer_identity,
    subject_did: str,
    claims: dict,
    disclosed: Optional[list[str]] = None,
    ttl: int = 86400,
) -> Attestation:
    """Create a signed attestation about another agent.

    Args:
        issuer_identity: Issuer's :class:`~skytale_sdk.identity.AgentIdentity`
            (must have a private key).
        subject_did: DID of the agent being attested.
        claims: Dict of claims (e.g. ``{"task_score": 0.95}``).
        disclosed: Which claim keys to disclose. Undisclosed claims
            are hashed (SD-JWT style). If ``None``, all are disclosed.
        ttl: Attestation lifetime in seconds (default 86400 = 1 day).

    Returns:
        Signed :class:`Attestation`.
    """
    disclosed_keys = disclosed if disclosed is not None else list(claims.keys())
    now = time.time()

    # Generate salts for each claim (SD-JWT pattern)
    salt_map = {}
    for key in claims:
        salt_map[key] = os.urandom(16).hex()

    att = Attestation(
        issuer=issuer_identity.did,
        subject=subject_did,
        claims=claims,
        disclosed_keys=disclosed_keys,
        issued_at=now,
        expires_at=now + ttl,
        salt_map=salt_map,
    )

    # Sign the attestation payload
    payload = _attestation_payload(att)
    att.signature = issuer_identity.sign(payload)

    return att


def verify_attestation(attestation: Attestation, issuer_public_key: bytes) -> bool:
    """Verify an attestation's signature.

    Args:
        attestation: The attestation to verify.
        issuer_public_key: Ed25519 public key of the claimed issuer.

    Returns:
        ``True`` if the signature is valid and the attestation hasn't expired.
    """
    if attestation.signature is None:
        return False

    if attestation.expires_at and time.time() > attestation.expires_at:
        return False

    from skytale_sdk.identity import AgentIdentity

    identity = AgentIdentity(
        did=attestation.issuer,
        public_key=issuer_public_key,
    )

    payload = _attestation_payload(attestation)
    return identity.verify(payload, attestation.signature)


def disclosed_claims(attestation: Attestation) -> dict:
    """Return only the disclosed claims from an attestation.

    Undisclosed claims are replaced with their SD-JWT hash
    (``SHA-256(salt || key || value)``).

    Args:
        attestation: The attestation to extract claims from.

    Returns:
        Dict with disclosed claims in plaintext, undisclosed as hashes.
    """
    result = {}
    for key, value in attestation.claims.items():
        if key in attestation.disclosed_keys:
            result[key] = value
        else:
            salt = attestation.salt_map.get(key, "")
            hash_input = f"{salt}{key}{json.dumps(value, sort_keys=True)}"
            result[key] = f"_sd:{hashlib.sha256(hash_input.encode()).hexdigest()}"
    return result


def serialize_attestation(attestation: Attestation) -> str:
    """Serialize an attestation to a JSON string for transmission.

    Only disclosed claims are included in plaintext. Undisclosed claims
    are included as SD-JWT hashes.

    Args:
        attestation: The attestation to serialize.

    Returns:
        JSON string suitable for sending over a Skytale channel.
    """
    import base64

    return json.dumps({
        "type": "attestation",
        "issuer": attestation.issuer,
        "subject": attestation.subject,
        "claims": disclosed_claims(attestation),
        "disclosed_keys": attestation.disclosed_keys,
        "issued_at": attestation.issued_at,
        "expires_at": attestation.expires_at,
        "signature": base64.b64encode(attestation.signature).decode()
        if attestation.signature
        else None,
    })


def _attestation_payload(attestation: Attestation) -> bytes:
    """Compute the canonical payload bytes for signing/verification."""
    canonical = json.dumps(
        {
            "issuer": attestation.issuer,
            "subject": attestation.subject,
            "claims": attestation.claims,
            "salt_map": attestation.salt_map,
            "issued_at": attestation.issued_at,
            "expires_at": attestation.expires_at,
        },
        sort_keys=True,
    )
    return canonical.encode()
