"""Offline verification helpers for AVP proof artifacts."""

from __future__ import annotations

import hashlib
import json
from datetime import datetime, timezone
from typing import Any, Iterable, Optional

import base58
import jcs
from nacl.exceptions import BadSignatureError
from nacl.signing import VerifyKey

from agentveil.delegation import DelegationInvalid, verify_delegation


class ProofVerificationError(ValueError):
    """Raised when a proof artifact fails signature or semantic verification."""


def _did_to_public_key(did: str) -> bytes:
    if not isinstance(did, str) or not did.startswith("did:key:z"):
        raise ProofVerificationError("signer DID must be did:key")
    try:
        decoded = base58.b58decode(did[len("did:key:z"):])
    except Exception as exc:
        raise ProofVerificationError("signer DID is not valid base58") from exc
    if len(decoded) < 2 or decoded[:2] != b"\xed\x01":
        raise ProofVerificationError("signer DID is not Ed25519 did:key")
    public_key = decoded[2:]
    if len(public_key) != 32:
        raise ProofVerificationError("signer DID has invalid Ed25519 public key")
    return public_key


def _receipt_dict(packet: Any) -> dict[str, Any]:
    if hasattr(packet, "to_dict"):
        packet = packet.to_dict()
    if not isinstance(packet, dict):
        raise ProofVerificationError("proof packet must be a dict-like object")
    return packet


def _sha256_text(text: str) -> str:
    return hashlib.sha256(text.encode("utf-8")).hexdigest()


def _canonical_delegation_hash(receipt: dict[str, Any]) -> str:
    return hashlib.sha256(jcs.canonicalize(receipt)).hexdigest()


def _historical_delegation_verify(receipt: dict[str, Any]) -> None:
    """Verify DelegationReceipt signature without rejecting historical expiry."""
    valid_from = receipt.get("validFrom")
    if isinstance(valid_from, str):
        try:
            now = datetime.strptime(valid_from, "%Y-%m-%dT%H:%M:%SZ").replace(
                tzinfo=timezone.utc,
            )
        except ValueError:
            now = None
    else:
        now = None
    try:
        verify_delegation(receipt, now=now)
    except DelegationInvalid as exc:
        raise ProofVerificationError(f"DelegationReceipt invalid: {exc.reason}") from exc


def verify_signed_jcs(
    receipt_jcs: str,
    expected_signer_did: Optional[str] = None,
) -> dict[str, Any]:
    """Verify one DataIntegrityProof / eddsa-jcs-2022 JSON artifact.

    This is a signature-level helper. It proves which DID signed the JCS text.
    AVP semantic checks, trust-anchor checks, and cross-receipt hash linkage are
    handled by ``verify_proof_packet``.
    """
    if not isinstance(receipt_jcs, str) or not receipt_jcs:
        raise ProofVerificationError("receipt_jcs must be a non-empty string")
    try:
        receipt = json.loads(receipt_jcs)
    except json.JSONDecodeError as exc:
        raise ProofVerificationError("receipt_jcs is not valid JSON") from exc
    if not isinstance(receipt, dict):
        raise ProofVerificationError("receipt_jcs must decode to a JSON object")

    proof = receipt.get("proof")
    if not isinstance(proof, dict):
        raise ProofVerificationError("receipt proof is missing")
    if proof.get("type") != "DataIntegrityProof":
        raise ProofVerificationError("proof.type must be DataIntegrityProof")
    if proof.get("cryptosuite") != "eddsa-jcs-2022":
        raise ProofVerificationError("proof.cryptosuite must be eddsa-jcs-2022")

    verification_method = proof.get("verificationMethod")
    if not isinstance(verification_method, str) or "#" not in verification_method:
        raise ProofVerificationError("proof.verificationMethod is invalid")
    signer_did = verification_method.split("#", 1)[0]
    if expected_signer_did is not None and signer_did != expected_signer_did:
        raise ProofVerificationError("receipt signer does not match expected signer")

    proof_value = proof.get("proofValue")
    if not isinstance(proof_value, str) or not proof_value.startswith("z"):
        raise ProofVerificationError("proof.proofValue must be multibase-z")
    try:
        signature = base58.b58decode(proof_value[1:])
    except Exception as exc:
        raise ProofVerificationError("proof.proofValue is not valid base58") from exc

    body = {key: value for key, value in receipt.items() if key != "proof"}
    try:
        VerifyKey(_did_to_public_key(signer_did)).verify(
            jcs.canonicalize(body),
            signature,
        )
    except BadSignatureError as exc:
        raise ProofVerificationError("receipt signature verification failed") from exc

    return {
        "valid": True,
        "body": body,
        "signer_did": signer_did,
        "schema_version": body.get("schema_version"),
        "digest": _sha256_text(receipt_jcs),
    }


def _require_backend_trust(
    verified: dict[str, Any],
    trusted_backend_signer_dids: set[str],
    label: str,
) -> None:
    if verified["signer_did"] not in trusted_backend_signer_dids:
        raise ProofVerificationError(f"{label} signer is not trusted")


def _normalize_trust_map(
    trusted_backend_signer_dids: Optional[Iterable[str]],
    trusted_decision_signer_dids: Optional[Iterable[str]],
    trusted_execution_signer_dids: Optional[Iterable[str]],
    trusted_human_approval_signer_dids: Optional[Iterable[str]],
) -> dict[str, set[str]]:
    fallback = {
        did for did in (trusted_backend_signer_dids or [])
        if isinstance(did, str)
    }
    decision = {
        did for did in (trusted_decision_signer_dids or fallback)
        if isinstance(did, str)
    }
    execution = {
        did for did in (trusted_execution_signer_dids or fallback)
        if isinstance(did, str)
    }
    approval = {
        did for did in (trusted_human_approval_signer_dids or fallback)
        if isinstance(did, str)
    }
    if not decision or not execution or not approval:
        raise ProofVerificationError("trusted signer DID(s) are required")
    return {
        "decision": decision,
        "execution": execution,
        "approval": approval,
    }


def _verify_optional_backend_receipt(
    packet: dict[str, Any],
    key: str,
    trusted_backend_signer_dids: set[str],
    label: str,
) -> Optional[dict[str, Any]]:
    receipt_jcs = packet.get(key)
    if receipt_jcs is None:
        return None
    verified = verify_signed_jcs(receipt_jcs)
    _require_backend_trust(verified, trusted_backend_signer_dids, label)
    return verified


def _assert_equal(actual: Any, expected: Any, label: str) -> None:
    if actual != expected:
        raise ProofVerificationError(f"{label} mismatch")


def _check_shared_intent(
    *,
    decision: dict[str, Any],
    other: dict[str, Any],
    other_label: str,
) -> None:
    required = other.get("schema_version") in {
        "execution_receipt/2",
        "human_approval_receipt/2",
    }
    for field in ("agent_did", "action", "resource", "environment"):
        other_field = (
            "requester_agent_did"
            if field == "agent_did" and other_label == "approval"
            else field
        )
        if other_field not in other:
            if required:
                raise ProofVerificationError(f"{other_label}.{other_field} is required")
            continue
        _assert_equal(
            other[other_field],
            decision.get(field),
            f"{other_label}.{other_field}",
        )


def _check_schema_version(body: dict[str, Any], allowed: Iterable[str], label: str) -> None:
    schema = body.get("schema_version")
    if schema is not None and schema not in set(allowed):
        raise ProofVerificationError(f"{label} schema_version is unsupported")


def _requires_field(body: dict[str, Any], current_schema: str, field: str) -> bool:
    return body.get("schema_version") == current_schema or field in body


def verify_proof_packet(
    packet: Any,
    trusted_backend_signer_dids: Optional[Iterable[str]] = None,
    *,
    trusted_decision_signer_dids: Optional[Iterable[str]] = None,
    trusted_execution_signer_dids: Optional[Iterable[str]] = None,
    trusted_human_approval_signer_dids: Optional[Iterable[str]] = None,
) -> dict[str, Any]:
    """Verify an AVP proof packet's signatures, trusted signers, and linkage."""
    trust = _normalize_trust_map(
        trusted_backend_signer_dids,
        trusted_decision_signer_dids,
        trusted_execution_signer_dids,
        trusted_human_approval_signer_dids,
    )

    data = _receipt_dict(packet)
    delegation_receipt = data.get("delegation_receipt")
    if not isinstance(delegation_receipt, dict):
        raise ProofVerificationError("proof packet missing DelegationReceipt")
    _historical_delegation_verify(delegation_receipt)
    delegation_hash = _canonical_delegation_hash(delegation_receipt)

    decision = _verify_optional_backend_receipt(
        data, "decision_receipt_jcs", trust["decision"], "DecisionReceipt",
    )
    approval = _verify_optional_backend_receipt(
        data, "approval_receipt_jcs", trust["approval"], "HumanApprovalReceipt",
    )
    execution = _verify_optional_backend_receipt(
        data, "execution_receipt_jcs", trust["execution"], "ExecutionReceipt",
    )

    if decision is None:
        raise ProofVerificationError("proof packet missing DecisionReceipt")
    decision_body = decision["body"]
    _check_schema_version(decision_body, {"decision_receipt/1", "decision_receipt/2"}, "DecisionReceipt")
    outcome_status = data.get("outcome_status")
    decision_value = decision_body.get("decision")
    if outcome_status == "executed" and execution is None:
        raise ProofVerificationError("executed proof packet missing ExecutionReceipt")
    if outcome_status == "blocked":
        if decision_value != "BLOCK":
            raise ProofVerificationError("blocked proof packet requires BLOCK DecisionReceipt")
        if execution is not None:
            raise ProofVerificationError("blocked proof packet must not include ExecutionReceipt")
    if outcome_status == "approval_required" and decision_value != "WAITING_FOR_HUMAN_APPROVAL":
        raise ProofVerificationError(
            "approval_required proof packet requires WAITING_FOR_HUMAN_APPROVAL DecisionReceipt"
        )
    _assert_equal(data.get("agent_did"), decision_body.get("agent_did"), "packet.agent_did")
    if data.get("audit_id") is not None:
        _assert_equal(data["audit_id"], decision_body.get("audit_id"), "packet.audit_id")
    if decision_body.get("delegation_receipt_hash") is not None:
        _assert_equal(
            decision_body["delegation_receipt_hash"],
            delegation_hash,
            "DecisionReceipt.delegation_receipt_hash",
        )

    if execution is not None:
        execution_body = execution["body"]
        _check_schema_version(execution_body, {"execution_receipt/1", "execution_receipt/2"}, "ExecutionReceipt")
        if _requires_field(execution_body, "execution_receipt/2", "decision_receipt_hash"):
            _assert_equal(
                execution_body.get("decision_receipt_hash"),
                decision["digest"],
                "ExecutionReceipt.decision_receipt_hash",
            )
        _assert_equal(
            execution_body.get("gate_audit_id"),
            decision_body.get("audit_id"),
            "ExecutionReceipt.gate_audit_id",
        )
        _check_shared_intent(
            decision=decision_body,
            other=execution_body,
            other_label="execution",
        )

    if approval is not None:
        approval_body = approval["body"]
        _check_schema_version(
            approval_body,
            {"human_approval_receipt/1", "human_approval_receipt/2"},
            "HumanApprovalReceipt",
        )
        if _requires_field(
            approval_body,
            "human_approval_receipt/2",
            "decision_receipt_hash",
        ):
            _assert_equal(
                approval_body.get("decision_receipt_hash"),
                decision["digest"],
                "HumanApprovalReceipt.decision_receipt_hash",
            )
        _assert_equal(
            approval_body.get("gate_audit_id"),
            decision_body.get("audit_id"),
            "HumanApprovalReceipt.gate_audit_id",
        )
        if approval_body.get("delegation_receipt_hash") is not None:
            _assert_equal(
                approval_body["delegation_receipt_hash"],
                delegation_hash,
                "HumanApprovalReceipt.delegation_receipt_hash",
            )
        _check_shared_intent(
            decision=decision_body,
            other=approval_body,
            other_label="approval",
        )

    approval_required = (
        data.get("outcome_status") == "approval_required"
        or (execution is not None and "approval_receipt_hash" in execution["body"])
        or data.get("approval") is not None
    )
    if approval_required and approval is None:
        raise ProofVerificationError("approval receipt is required for approval path")

    if execution is not None and approval is not None:
        _assert_equal(
            execution["body"].get("approval_receipt_hash"),
            approval["digest"],
            "ExecutionReceipt.approval_receipt_hash",
        )
        if approval["body"].get("decision") != "APPROVED":
            raise ProofVerificationError("approval receipt is not APPROVED")

    return {
        "valid": True,
        "decision_receipt": decision,
        "approval_receipt": approval,
        "execution_receipt": execution,
        "delegation_receipt_hash": delegation_hash,
    }


__all__ = [
    "ProofVerificationError",
    "verify_signed_jcs",
    "verify_proof_packet",
]
