"""PlanCert and PlanProof helpers.

Lightweight, dependency-free PlanCert representation and HMAC-based signing
for constrained environments. This is intentionally small: in production one
would replace the HMAC signer with a proper asymmetric signature scheme
and a compact binary encoding (CBOR) for wire efficiency.

Provides:
- PlanCert: declarative certificate describing plan invariants and provenance
- PlanProof: verifier result (accept/reject + score + optional counterexample)
- helpers: sign_plan_cert, verify_plan_cert, make_plan_proof
"""
from __future__ import annotations

import json
import hmac
import hashlib
from dataclasses import dataclass, asdict
from typing import Any, Dict, Optional


def _canonical_json(obj: Any) -> str:
    # Deterministic serialization used for signatures
    return json.dumps(obj, sort_keys=True, separators=(",", ":"))


@dataclass
class PlanCert:
    plan_id: str
    preconditions: Dict[str, Any]
    postconditions: Dict[str, Any]
    invariants: Dict[str, Any]
    worst_case_energy: Optional[float] = None
    wcet_ms: Optional[int] = None
    provenance: Dict[str, Any] = None
    signature: Optional[str] = None  # hex HMAC-SHA256 signature when signed

    def to_dict(self) -> Dict[str, Any]:
        d = asdict(self)
        # don't include signature when computing signatures
        return {k: v for k, v in d.items() if k != "signature"}

    def to_json(self) -> str:
        return _canonical_json(self.to_dict())


@dataclass
class PlanProof:
    plan_id: str
    accepted: bool
    score: float
    counterexample: Optional[Dict[str, Any]] = None
    verifier: Optional[str] = None
    signature: Optional[str] = None

    def to_dict(self) -> Dict[str, Any]:
        # Exclude signature field when producing the canonical dict used for
        # signing/verification to avoid the signature covering itself.
        d = asdict(self)
        return {k: v for k, v in d.items() if k != "signature"}

    def to_json(self) -> str:
        return _canonical_json(self.to_dict())


def sign_plan_cert(cert: PlanCert, key: bytes) -> PlanCert:
    """Sign PlanCert with HMAC-SHA256 using symmetric key.

    The signature covers all fields except `signature` itself.
    """
    data = cert.to_json().encode()
    sig = hmac.new(key, data, hashlib.sha256).hexdigest()
    cert.signature = sig
    return cert


def verify_plan_cert(cert: PlanCert, key: bytes) -> bool:
    if not cert.signature:
        return False
    expected = hmac.new(key, cert.to_json().encode(), hashlib.sha256).hexdigest()
    # Use hmac.compare_digest for timing-safe comparison
    return hmac.compare_digest(expected, cert.signature)


def make_plan_proof(cert: PlanCert, accepted: bool, score: float, verifier: str, key: Optional[bytes] = None, counterexample: Optional[Dict[str, Any]] = None) -> PlanProof:
    proof = PlanProof(plan_id=cert.plan_id, accepted=accepted, score=score, counterexample=counterexample, verifier=verifier)
    if key is not None:
        proof.signature = hmac.new(key, proof.to_json().encode(), hashlib.sha256).hexdigest()
    return proof


def verify_plan_proof(proof: PlanProof, key: bytes) -> bool:
    if not proof.signature:
        return False
    expected = hmac.new(key, proof.to_json().encode(), hashlib.sha256).hexdigest()
    return hmac.compare_digest(expected, proof.signature)