"""PlanCert and PlanProof utilities.

Provides small, verifiable plan certificate schema and a lightweight verifier
that performs basic syntactic and semantic checks using the existing safety
predicate evaluator. This is intentionally tiny and suitable for edge use.
"""
from __future__ import annotations

from typing import Any, Dict, List, Optional, Tuple
from pydantic import BaseModel

from .safety import evaluate_predicate, UnsafeExpression


class PlanCert(BaseModel):
    plan_id: str
    planner_version: Optional[str] = None
    preconditions: List[str] = []
    postconditions: List[str] = []
    invariants: Dict[str, str] = {}
    worst_case_energy: Optional[float] = None
    max_actuator_exposure: Optional[float] = None
    environment_hash: Optional[str] = None


class PlanProof(BaseModel):
    plan_id: str
    accepted: bool
    score: float
    reasons: List[str] = []
    counterexample: Optional[Dict[str, Any]] = None


def verify_plan(cert: PlanCert, initial_state: Dict[str, Any], energy_budget: Optional[float] = None) -> PlanProof:
    """Verify a PlanCert against an initial state and optional energy budget.

    Checks performed:
    - All preconditions evaluate to True on the provided initial_state.
    - All invariants evaluate to True on the provided initial_state.
    - If energy_budget provided and cert.worst_case_energy is set, ensure it's <= budget.

    Returns a PlanProof with accepted flag and reasons. This is a compact,
    deterministic verifier; probabilistic adversarial sampling is outside its
    scope.
    """
    reasons: List[str] = []
    score = 1.0

    # Check preconditions
    for i, p in enumerate(cert.preconditions):
        try:
            ok = evaluate_predicate(p, initial_state)
        except UnsafeExpression as e:
            reasons.append(f"precondition[{i}] unsafe: {e}")
            ok = False
        if not ok:
            reasons.append(f"precondition[{i}] failed: {p}")
            score -= 0.4

    # Check invariants
    for name, inv in cert.invariants.items():
        try:
            ok = evaluate_predicate(inv, initial_state)
        except UnsafeExpression as e:
            reasons.append(f"invariant[{name}] unsafe: {e}")
            ok = False
        if not ok:
            reasons.append(f"invariant[{name}] failed: {inv}")
            score -= 0.3

    # Energy budget check
    if energy_budget is not None and cert.worst_case_energy is not None:
        try:
            if cert.worst_case_energy > energy_budget:
                reasons.append(f"worst_case_energy {cert.worst_case_energy} > budget {energy_budget}")
                score -= 0.5
        except Exception:
            reasons.append("invalid energy comparison")
            score -= 0.5

    accepted = score > 0.0 and len([r for r in reasons if r.startswith("precondition") or r.startswith("invariant")]) == 0

    # If rejected, include a minimal counterexample: the parts of state that caused failures
    counterexample = None
    if not accepted:
        counterexample = {k: initial_state.get(k) for k in sorted(initial_state.keys())[:10]}

    # Clamp score
    if score < 0:
        score = 0.0
    if score > 1.0:
        score = 1.0

    return PlanProof(plan_id=cert.plan_id, accepted=accepted, score=score, reasons=reasons, counterexample=counterexample)