"""
PolicyEngine — applies static rules + learned heuristics + live overrides
to scheduler decisions.

Constructor accepts an optional `db` (anything with an `events` dict-like
attribute). If db is None, the heuristics layer is skipped.
"""


class PolicyEngine:

    def __init__(self, db=None):
        self.db = db
        self.rules = []
        self.overrides = {}

    # -----------------------------
    # REGISTER RULE
    # -----------------------------
    def add_rule(self, rule):
        self.rules.append(rule)

    # -----------------------------
    # SET LIVE OVERRIDE
    # -----------------------------
    def set_override(self, key, value):
        self.overrides[key] = value

    # -----------------------------
    # CLEAR OVERRIDE
    # -----------------------------
    def clear_override(self, key):
        if key in self.overrides:
            del self.overrides[key]

    # -----------------------------
    # APPLY POLICY TO NODE SELECTION
    # -----------------------------
    def evaluate(self, action, target, node):
        score_modifier = 0
        for rule in self.rules:
            try:
                score_modifier += rule.apply(action, target, node)
            except Exception:
                pass
        score_modifier += self._heuristics(action, target, node)
        score_modifier += self._overrides(action, target, node)
        return score_modifier

    # -----------------------------
    # HEURISTICS (LEARNING LAYER)
    # -----------------------------
    def _heuristics(self, action, target, node):
        if self.db is None:
            return 0
        if not hasattr(self.db, "events"):
            return 0
        try:
            history = self.db.events
        except Exception:
            return 0

        score = 0
        success_count = 0
        # history may be dict-of-lists or a flat list
        if isinstance(history, dict):
            iterable = (e for sess in history.values() for e in (sess or []))
        else:
            iterable = history or []

        for event in iterable:
            if not isinstance(event, dict):
                continue
            if event.get("type") in ("action_end", "task_update"):
                d = event.get("data", event)
                if (d.get("action") == action.get("name")
                    and d.get("state") == "done"
                    and d.get("node") == node.get("name")):
                    success_count += 1
        score += success_count * 2
        return score

    # -----------------------------
    # LIVE OVERRIDES
    # -----------------------------
    def _overrides(self, action, target, node):
        score = 0
        forced = self.overrides.get("force_node")
        if forced and node.get("name") == forced:
            score += 1000
        avoid = self.overrides.get("avoid_node")
        if avoid and node.get("name") == avoid:
            score -= 1000
        return score