"""Reusable vote-proof primitives for vote-mcp vote-submit clients.

This module is client-side helper code for constructing the exact canonical
vote-proof statement that the server verifies.
"""

from __future__ import annotations

import base64
import hashlib
import json
import re
from datetime import datetime, timezone as dt_timezone

from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey

VOTE_PROOF_SCHEME = "ed25519_statement_v1"
VOTE_PROOF_STATEMENT_PREFIX = "vote-mcp:signature-statement:v1\n"

_SEED_HEX_RE = re.compile(r"^[0-9a-f]{64}$")
_SHA256_RE = re.compile(r"^sha256:[0-9a-f]{64}$")
_RFC3339_UTC_Z_RE = re.compile(r"^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}(?:\.\d{1,6})?Z$")


def canonical_json_bytes(payload: object) -> bytes:
    """Deterministically serialize JSON using server-compatible canonical settings."""
    try:
        return json.dumps(
            payload,
            separators=(",", ":"),
            sort_keys=True,
            allow_nan=False,
        ).encode("utf-8")
    except (TypeError, ValueError):
        raise ValueError("payload is not canonically JSON-serializable.") from None


def canonical_response_hash(response_payload: object) -> str:
    """Build `sha256:<hex>` over canonical JSON bytes for `response`."""
    if not isinstance(response_payload, dict):
        raise ValueError("response_payload must be a JSON object.")
    digest = hashlib.sha256(canonical_json_bytes(response_payload)).hexdigest()
    return f"sha256:{digest}"


def private_key_from_seed_hex(seed_hex: str) -> Ed25519PrivateKey:
    """Reconstruct Ed25519 private key from a 32-byte lowercase hex seed."""
    if not isinstance(seed_hex, str) or _SEED_HEX_RE.fullmatch(seed_hex) is None:
        raise ValueError("seed_hex must be 32-byte lowercase hex (64 chars).")
    return Ed25519PrivateKey.from_private_bytes(bytes.fromhex(seed_hex))


def public_key_base64url_from_private_key(private_key: Ed25519PrivateKey) -> str:
    """Encode Ed25519 public key as canonical unpadded base64url."""
    raw = private_key.public_key().public_bytes(
        encoding=serialization.Encoding.Raw,
        format=serialization.PublicFormat.Raw,
    )
    return base64.urlsafe_b64encode(raw).decode("ascii").rstrip("=")


def format_rfc3339_utc_z(value: datetime) -> str:
    """Normalize a timezone-aware datetime to RFC3339 UTC with `Z` suffix."""
    if not isinstance(value, datetime):
        raise ValueError("value must be a datetime instance.")
    if value.tzinfo is None or value.utcoffset() is None:
        raise ValueError("value must be timezone-aware.")
    return value.astimezone(dt_timezone.utc).isoformat().replace("+00:00", "Z")


def _normalize_issued_at(issued_at: str) -> str:
    if not isinstance(issued_at, str) or _RFC3339_UTC_Z_RE.fullmatch(issued_at) is None:
        raise ValueError("issued_at must be RFC3339 UTC with 'Z' suffix.")
    try:
        parsed = datetime.fromisoformat(issued_at.replace("Z", "+00:00"))
    except ValueError:
        raise ValueError("issued_at must be RFC3339 UTC with 'Z' suffix.") from None
    if parsed.tzinfo is None or parsed.utcoffset() is None:
        raise ValueError("issued_at must be timezone-aware UTC.")
    return format_rfc3339_utc_z(parsed)


def build_vote_proof_statement_claims(
    *,
    poll_id: str,
    vote_id: str,
    response_hash: str,
    challenge_nonce: str,
    issued_at: str,
) -> dict[str, str]:
    """Build canonical vote-proof claims in the exact schema verified by server."""
    if not isinstance(poll_id, str) or not poll_id:
        raise ValueError("poll_id must be a non-empty string.")
    if not isinstance(vote_id, str) or not vote_id:
        raise ValueError("vote_id must be a non-empty string.")
    if not isinstance(response_hash, str) or _SHA256_RE.fullmatch(response_hash) is None:
        raise ValueError("response_hash must match `sha256:<64 lowercase hex>`.")
    if not isinstance(challenge_nonce, str) or not challenge_nonce or any(
        ch.isspace() for ch in challenge_nonce
    ):
        raise ValueError("challenge_nonce must be a non-empty token string.")

    normalized_issued_at = _normalize_issued_at(issued_at)

    return {
        "domain": "vote-mcp",
        "kind": "vote_proof",
        "version": "v1",
        "poll_id": poll_id,
        "vote_id": vote_id,
        "response_hash": response_hash,
        "challenge_nonce": challenge_nonce,
        "issued_at": normalized_issued_at,
    }


def build_vote_proof_statement_bytes(
    *,
    poll_id: str,
    vote_id: str,
    response_hash: str,
    challenge_nonce: str,
    issued_at: str,
) -> bytes:
    """Build canonical statement bytes signed by `proof.signature`."""
    claims = build_vote_proof_statement_claims(
        poll_id=poll_id,
        vote_id=vote_id,
        response_hash=response_hash,
        challenge_nonce=challenge_nonce,
        issued_at=issued_at,
    )
    return VOTE_PROOF_STATEMENT_PREFIX.encode("utf-8") + canonical_json_bytes(claims)


def build_vote_proof(
    *,
    private_key_seed_hex: str,
    poll_id: str,
    vote_id: str,
    response_hash: str,
    challenge_nonce: str,
    issued_at: str,
) -> dict[str, str]:
    """Assemble a complete vote-proof object for `POST /api/v1/polls/{poll_id}/vote`."""
    private_key = private_key_from_seed_hex(private_key_seed_hex)
    statement = build_vote_proof_statement_bytes(
        poll_id=poll_id,
        vote_id=vote_id,
        response_hash=response_hash,
        challenge_nonce=challenge_nonce,
        issued_at=issued_at,
    )
    signature = base64.urlsafe_b64encode(private_key.sign(statement)).decode("ascii").rstrip("=")
    return {
        "scheme": VOTE_PROOF_SCHEME,
        "public_key": public_key_base64url_from_private_key(private_key),
        "issued_at": _normalize_issued_at(issued_at),
        "signature": signature,
    }