"""PlanCert schema and lightweight verifier.

Provides PlanCert and PlanProof models and a small, safe expression evaluator used by
verify_plan to check preconditions and simple invariants. This is intentionally small
and dependency-free beyond pydantic which the project already depends on.
"""
from __future__ import annotations

from typing import Dict, Any, List, Optional
import ast

from pydantic import BaseModel, Field


class PlanCert(BaseModel):
    plan_id: str
    preconditions: List[str] = Field(default_factory=list)
    invariants: Dict[str, str] = Field(default_factory=dict)
    worst_case_energy: float = 0.0
    provenance: Dict[str, Any] = Field(default_factory=dict)


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


class _SafeEvalError(Exception):
    pass


def _safe_eval(expr: str, state: Dict[str, Any]) -> bool:
    """Evaluate simple boolean/comparison expressions against a state dict.

    Allowed expressions: comparisons (>, >=, <, <=, ==, !=), boolean and/or, numeric
    literals, names that map to keys in `state` and simple arithmetic (+,-,*,/).
    This uses ast to prevent arbitrary code execution.
    """
    node = ast.parse(expr, mode="eval")

    def _eval(n):
        if isinstance(n, ast.Expression):
            return _eval(n.body)
        if isinstance(n, ast.BoolOp):
            if isinstance(n.op, ast.And):
                return all(_eval(v) for v in n.values)
            if isinstance(n.op, ast.Or):
                return any(_eval(v) for v in n.values)
            raise _SafeEvalError("unsupported boolean op")
        if isinstance(n, ast.Compare):
            left = _eval(n.left)
            for op, comp in zip(n.ops, n.comparators):
                right = _eval(comp)
                if isinstance(op, ast.Gt):
                    ok = left > right
                elif isinstance(op, ast.GtE):
                    ok = left >= right
                elif isinstance(op, ast.Lt):
                    ok = left < right
                elif isinstance(op, ast.LtE):
                    ok = left <= right
                elif isinstance(op, ast.Eq):
                    ok = left == right
                elif isinstance(op, ast.NotEq):
                    ok = left != right
                else:
                    raise _SafeEvalError("unsupported comparison")
                if not ok:
                    return False
                left = right
            return True
        if isinstance(n, ast.BinOp):
            l = _eval(n.left)
            r = _eval(n.right)
            if isinstance(n.op, ast.Add):
                return l + r
            if isinstance(n.op, ast.Sub):
                return l - r
            if isinstance(n.op, ast.Mult):
                return l * r
            if isinstance(n.op, ast.Div):
                return l / r
            raise _SafeEvalError("unsupported binary operator")
        if isinstance(n, ast.UnaryOp):
            if isinstance(n.op, ast.USub):
                return -_eval(n.operand)
            raise _SafeEvalError("unsupported unary op")
        if isinstance(n, ast.Name):
            if n.id in state:
                return state[n.id]
            raise _SafeEvalError(f"unknown name: {n.id}")
        if isinstance(n, ast.Constant):
            return n.value
        if isinstance(n, ast.Num):  # type: ignore
            return n.n  # type: ignore
        raise _SafeEvalError(f"unsupported AST node: {type(n)}")

    return bool(_eval(node))


def verify_plan(cert: PlanCert, state: Dict[str, Any], energy_budget: float) -> PlanProof:
    """Run a lightweight verification of the PlanCert against a given state and energy budget.

    Checks:
    - All preconditions evaluate to True in `state`.
    - All invariants evaluate to True in `state`.
    - worst_case_energy <= energy_budget

    Returns a PlanProof with accepted bool, a simple confidence score (0..1), and reasons.
    """
    reasons: List[str] = []

    # Check preconditions
    for p in cert.preconditions:
        try:
            ok = _safe_eval(p, state)
        except Exception as e:
            reasons.append(f"precondition parse error: {p} -> {e}")
            ok = False
        if not ok:
            reasons.append(f"precondition failed: {p}")

    # Check invariants
    for name, inv in cert.invariants.items():
        try:
            ok = _safe_eval(inv, state)
        except Exception as e:
            reasons.append(f"invariant parse error ({name}): {inv} -> {e}")
            ok = False
        if not ok:
            reasons.append(f"invariant failed ({name}): {inv}")

    # Energy check
    if cert.worst_case_energy is not None and cert.worst_case_energy > energy_budget:
        reasons.append(
            f"worst_case_energy {cert.worst_case_energy} exceeds budget {energy_budget}"
        )

    accepted = len(reasons) == 0

    # Simple score: start at 1.0, subtract 0.25 per reason, clamp to [0,1]
    score = max(0.0, 1.0 - 0.25 * len(reasons))

    proof = PlanProof(accepted=accepted, score=score, reasons=reasons or [])
    return proof


__all__ = ["PlanCert", "PlanProof", "verify_plan"]