tiny_trigger/automation.py

from __future__ import annotations

import json
from collections import defaultdict
from math import hypot
from statistics import median
from typing import Any, Literal

from pydantic import AliasChoices, BaseModel, ConfigDict, Field, ValidationError, model_validator

from .models import ActionEvent, Detection, FrameSample


def _label(value: str) -> str:
    return value.strip().lower()


def _dedupe_labels(labels: list[str]) -> list[str]:
    seen: set[str] = set()
    result: list[str] = []
    for label in labels:
        normalized = " ".join(label.strip().split())
        key = normalized.lower()
        if normalized and key not in seen:
            seen.add(key)
            result.append(normalized)
    return result


class PresentCondition(BaseModel):
    label: str
    min_count: int = Field(default=1, ge=1)


class CountCondition(BaseModel):
    label: str
    minimum: int = Field(default=1, ge=1, validation_alias=AliasChoices("min", "minimum", "min_count"))


class NearCondition(BaseModel):
    a: str
    b: str
    max_gap_percent: float = Field(
        default=16.0,
        ge=0.0,
        validation_alias=AliasChoices("max_gap_percent", "max_distance"),
    )

    @model_validator(mode="before")
    @classmethod
    def migrate_max_distance(cls, data: Any) -> Any:
        if isinstance(data, dict) and "max_distance" in data and "max_gap_percent" not in data:
            value = data["max_distance"]
            if isinstance(value, int | float) and value <= 1:
                return {**data, "max_gap_percent": value * 100.0}
        return data


class FarCondition(BaseModel):
    a: str
    b: str
    min_gap_percent: float = Field(default=25.0, ge=0.0)


class MovingCondition(BaseModel):
    label: str
    min_displacement_ratio: float = Field(default=0.15, ge=0.0)
    window_frames: int = Field(default=3, ge=3)
    max_missing_frames: int = Field(default=1, ge=0)

    @model_validator(mode="before")
    @classmethod
    def migrate_short_window(cls, data: Any) -> Any:
        if isinstance(data, dict) and data.get("window_frames") is not None:
            try:
                window_frames = int(data["window_frames"])
            except (TypeError, ValueError):
                return data
            if window_frames < 3:
                return {**data, "window_frames": 3}
        return data


class CooldownCondition(BaseModel):
    key: str | None = None
    seconds: float | None = Field(default=None, ge=0.0)
    minutes: float | None = Field(default=None, ge=0.0)

    @property
    def duration_seconds(self) -> float:
        if self.seconds is not None:
            return self.seconds
        if self.minutes is not None:
            return self.minutes * 60.0
        return 0.0


class ConditionBlock(BaseModel):
    present: PresentCondition | None = None
    count: CountCondition | None = None
    near: NearCondition | None = None
    far: FarCondition | None = None
    moving: MovingCondition | None = None
    cooldown: CooldownCondition | None = None

    @model_validator(mode="after")
    def exactly_one_condition(self) -> "ConditionBlock":
        selected = [self.present, self.count, self.near, self.far, self.moving, self.cooldown]
        if sum(item is not None for item in selected) != 1:
            raise ValueError("Each condition block must contain exactly one condition.")
        return self


class WhenClause(BaseModel):
    model_config = ConfigDict(populate_by_name=True)

    all_conditions: list[ConditionBlock] = Field(default_factory=list, alias="all")
    any_conditions: list[ConditionBlock] = Field(default_factory=list, alias="any")

    @model_validator(mode="after")
    def has_conditions(self) -> "WhenClause":
        if not self.all_conditions and not self.any_conditions:
            raise ValueError("A rule needs at least one condition.")
        return self


class ActionSpec(BaseModel):
    type: Literal["simulate", "webhook"] = "simulate"
    name: str
    url: str | None = None
    payload: dict[str, Any] = Field(default_factory=dict)


class TriggerClause(BaseModel):
    on: Literal["while", "enter", "exit", "change"] = "enter"


class ActionSet(BaseModel):
    model_config = ConfigDict(populate_by_name=True)

    enter: list[ActionSpec] = Field(default_factory=list)
    exit: list[ActionSpec] = Field(default_factory=list)
    while_actions: list[ActionSpec] = Field(default_factory=list, alias="while")


class GateClause(BaseModel):
    enabled: bool = True
    cooldown: CooldownCondition | None = None


class AutomationRule(BaseModel):
    name: str
    when: WhenClause
    trigger: TriggerClause = Field(default_factory=TriggerClause)
    gate: GateClause = Field(default_factory=GateClause)
    then: list[ActionSpec] | ActionSet

    @model_validator(mode="after")
    def has_actions_for_trigger(self) -> "AutomationRule":
        if isinstance(self.then, list):
            if not self.then:
                raise ValueError("A rule needs at least one action.")
            if self.trigger.on == "change":
                raise ValueError("change triggers need then.enter/then.exit actions.")
            return self

        actions = _actions_for_trigger(self.then, self.trigger.on)
        if not actions:
            raise ValueError(f"A rule with trigger.on={self.trigger.on!r} needs matching actions.")
        return self


class AutomationDocument(BaseModel):
    rules: list[AutomationRule] = Field(default_factory=list)


def load_automation_text(text: str) -> AutomationDocument:
    """Load a JSON/YAML automation document and validate it."""
    raw = text.strip()
    if not raw:
        raise ValueError("Rules are empty.")

    try:
        data = json.loads(raw)
    except json.JSONDecodeError:
        try:
            import yaml
        except ImportError as exc:  # pragma: no cover - dependency guard
            raise RuntimeError("Install PyYAML to load YAML automation rules.") from exc
        data = yaml.safe_load(raw)

    if isinstance(data, list):
        data = {"rules": data}
    if isinstance(data, dict) and "rules" not in data and "name" in data:
        data = {"rules": [data]}

    try:
        return AutomationDocument.model_validate(data)
    except ValidationError:
        raise
    except Exception as exc:
        raise ValueError(f"Invalid automation document: {exc}") from exc


def automation_schema() -> dict[str, Any]:
    return AutomationDocument.model_json_schema()


def rule_labels(rule: AutomationRule) -> list[str]:
    labels: list[str] = []
    for condition in [*rule.when.all_conditions, *rule.when.any_conditions]:
        if condition.present:
            labels.append(condition.present.label)
        if condition.count:
            labels.append(condition.count.label)
        if condition.near:
            labels.extend([condition.near.a, condition.near.b])
        if condition.far:
            labels.extend([condition.far.a, condition.far.b])
        if condition.moving:
            labels.append(condition.moving.label)
    return _dedupe_labels(labels)


def document_labels(document: AutomationDocument, *, enabled_only: bool = True) -> list[str]:
    labels: list[str] = []
    for rule in document.rules:
        if enabled_only and not rule.gate.enabled:
            continue
        labels.extend(rule_labels(rule))
    return _dedupe_labels(labels)


class RuleEngine:
    def __init__(
        self,
        rules: list[AutomationRule],
        last_fired: dict[str, float] | None = None,
        last_matched: dict[str, bool] | None = None,
    ) -> None:
        self.rules = rules
        self.last_fired: dict[str, float] = dict(last_fired or {})
        self.last_matched: dict[str, bool] = dict(last_matched or {})
        self.track_history: dict[int, list[tuple[int, tuple[float, float], float]]] = {}
        self.moving_track_last_seen: dict[tuple[str, int], int] = {}

    def evaluate_frame(
        self,
        detections: list[Detection],
        *,
        frame_index: int,
        timestamp_sec: float,
    ) -> list[ActionEvent]:
        events: list[ActionEvent] = []
        self._update_track_history(detections)
        for rule in self.rules:
            if not self._gate_allows(rule, timestamp_sec):
                self.last_matched[rule.name] = False
                continue
            matched = self._rule_matches(rule, detections, frame_index, timestamp_sec)
            previous = self.last_matched.get(rule.name, False)
            edge = _trigger_edge(previous=previous, matched=matched)
            self.last_matched[rule.name] = matched
            actions = _actions_to_fire(rule, edge)
            if not matched and not actions:
                continue

            self._mark_cooldowns(rule, timestamp_sec)
            for action in actions:
                payload = {
                    "rule": rule.name,
                    "action": action.name,
                    "trigger": edge,
                    "frame_index": frame_index,
                    "timestamp_sec": timestamp_sec,
                    "detections": [item.model_dump(mode="json") for item in detections],
                }
                payload.update(action.payload)
                events.append(
                    ActionEvent(
                        rule=rule.name,
                        action=action.name,
                        type=action.type,
                        frame_index=frame_index,
                        timestamp_sec=timestamp_sec,
                        url=action.url,
                        payload=payload,
                    )
                )
        return events

    def _rule_matches(
        self,
        rule: AutomationRule,
        detections: list[Detection],
        frame_index: int,
        timestamp_sec: float,
    ) -> bool:
        all_ok = all(
            self._condition_matches(condition, detections, frame_index, rule.name, timestamp_sec)
            for condition in rule.when.all_conditions
        )
        any_ok = True
        if rule.when.any_conditions:
            any_ok = any(
                self._condition_matches(condition, detections, frame_index, rule.name, timestamp_sec)
                for condition in rule.when.any_conditions
            )
        return all_ok and any_ok

    def _gate_allows(self, rule: AutomationRule, timestamp_sec: float) -> bool:
        if not rule.gate.enabled:
            return False
        if not rule.gate.cooldown:
            return True
        return self._cooldown_allows(rule.gate.cooldown, rule.name, timestamp_sec)

    def _condition_matches(
        self,
        condition: ConditionBlock,
        detections: list[Detection],
        frame_index: int,
        rule_name: str,
        timestamp_sec: float,
    ) -> bool:
        by_label = _group_by_label(detections)

        if condition.present:
            return len(by_label[_label(condition.present.label)]) >= condition.present.min_count

        if condition.count:
            return len(by_label[_label(condition.count.label)]) >= condition.count.minimum

        if condition.near:
            left = by_label[_label(condition.near.a)]
            right = by_label[_label(condition.near.b)]
            if not left or not right:
                return False
            return _min_box_gap_percent(left, right) <= condition.near.max_gap_percent

        if condition.far:
            left = by_label[_label(condition.far.a)]
            right = by_label[_label(condition.far.b)]
            if not left or not right:
                return False
            return _min_box_gap_percent(left, right) >= condition.far.min_gap_percent

        if condition.moving:
            return self._moving_matches(condition.moving, by_label, frame_index)

        if condition.cooldown:
            return self._cooldown_allows(condition.cooldown, rule_name, timestamp_sec)

        return False

    def _moving_matches(
        self,
        condition: MovingCondition,
        by_label: dict[str, list[Detection]],
        frame_index: int,
    ) -> bool:
        label = _label(condition.label)
        label_detections = by_label[label]
        for detection in label_detections:
            if detection.track_id is None:
                continue
            history = self.track_history.get(detection.track_id, [])
            if len(history) < condition.window_frames:
                continue
            window = history[-condition.window_frames :]
            # Smooth out box jitter by averaging each half of the window, then
            # measure displacement relative to the object's own size (its median
            # box diagonal) so the threshold is distance-invariant and a flickering
            # box on a stationary object stays well under it.
            half = condition.window_frames // 2
            early = _mean_point([point for _frame, point, _size in window[:half]])
            late = _mean_point([point for _frame, point, _size in window[-half:]])
            size = median(size for _frame, _point, size in window)
            if size <= 0:
                continue
            displacement = hypot(late[0] - early[0], late[1] - early[1])
            if displacement / size >= condition.min_displacement_ratio:
                self.moving_track_last_seen[(label, detection.track_id)] = detection.frame_index
                return True
        if label_detections:
            return False
        return any(
            last_seen_frame <= frame_index
            and frame_index - last_seen_frame <= condition.max_missing_frames
            for (track_label, _track_id), last_seen_frame in self.moving_track_last_seen.items()
            if track_label == label
        )

    def _update_track_history(self, detections: list[Detection]) -> None:
        for detection in detections:
            if detection.track_id is None:
                continue
            history = self.track_history.setdefault(detection.track_id, [])
            history.append(
                (detection.frame_index, _foot_point_percent(detection), _box_diagonal_percent(detection))
            )
            del history[:-10]

    def _mark_cooldowns(self, rule: AutomationRule, timestamp_sec: float) -> None:
        if rule.gate.cooldown:
            self.last_fired[rule.gate.cooldown.key or rule.name] = timestamp_sec
        for condition in [*rule.when.all_conditions, *rule.when.any_conditions]:
            if condition.cooldown:
                self.last_fired[condition.cooldown.key or rule.name] = timestamp_sec

    def _cooldown_allows(self, cooldown: CooldownCondition, rule_name: str, timestamp_sec: float) -> bool:
        key = cooldown.key or rule_name
        last = self.last_fired.get(key)
        return last is None or (timestamp_sec - last) >= cooldown.duration_seconds


def evaluate_video_detections(
    rules: list[AutomationRule],
    detections: list[Detection],
    *,
    frames: list[FrameSample] | None = None,
    last_fired: dict[str, float] | None = None,
) -> tuple[list[ActionEvent], dict[str, float]]:
    engine = RuleEngine(rules, last_fired=last_fired)
    events: list[ActionEvent] = []
    grouped: dict[tuple[int, float], list[Detection]] = defaultdict(list)

    for detection in detections:
        grouped[(detection.frame_index, detection.timestamp_sec)].append(detection)
    if frames:
        for frame in frames:
            grouped.setdefault((frame.frame_index, frame.timestamp_sec), [])

    for (frame_index, timestamp_sec), frame_detections in sorted(grouped.items()):
        events.extend(
            engine.evaluate_frame(
                frame_detections,
                frame_index=frame_index,
                timestamp_sec=timestamp_sec,
            )
        )
    return events, dict(engine.last_fired)


def _group_by_label(detections: list[Detection]) -> dict[str, list[Detection]]:
    grouped: dict[str, list[Detection]] = defaultdict(list)
    for detection in detections:
        grouped[_label(detection.label)].append(detection)
    return grouped


def _min_box_gap_percent(left: list[Detection], right: list[Detection]) -> float:
    best = float("inf")
    for left_detection in left:
        for right_detection in right:
            best = min(best, _box_gap_percent(left_detection, right_detection))
    return best


def _box_gap_percent(left: Detection, right: Detection) -> float:
    ax1, ay1, ax2, ay2 = left.bbox_xyxy_norm
    bx1, by1, bx2, by2 = right.bbox_xyxy_norm
    gap_x = max(0.0, max(bx1 - ax2, ax1 - bx2))
    gap_y = max(0.0, max(by1 - ay2, ay1 - by2))
    return max(gap_x, gap_y) * 100.0


def _foot_point_percent(detection: Detection) -> tuple[float, float]:
    # Bottom-center ("foot point"): far steadier than the box center for a
    # standing/grounded object, whose top edge flickers with arms/hair/occlusion.
    x1, _y1, x2, y2 = detection.bbox_xyxy_norm
    return ((x1 + x2) * 50.0, y2 * 100.0)


def _box_diagonal_percent(detection: Detection) -> float:
    x1, y1, x2, y2 = detection.bbox_xyxy_norm
    return hypot((x2 - x1) * 100.0, (y2 - y1) * 100.0)


def _mean_point(points: list[tuple[float, float]]) -> tuple[float, float]:
    count = len(points)
    return (sum(point[0] for point in points) / count, sum(point[1] for point in points) / count)


def _trigger_edge(*, previous: bool, matched: bool) -> Literal["enter", "exit", "while", "none"]:
    if matched and not previous:
        return "enter"
    if not matched and previous:
        return "exit"
    if matched:
        return "while"
    return "none"


def _actions_to_fire(rule: AutomationRule, edge: str) -> list[ActionSpec]:
    trigger = rule.trigger.on
    if trigger == "while":
        return _actions_for_trigger(rule.then, "while") if edge in {"enter", "while"} else []
    if trigger == "enter":
        return _actions_for_trigger(rule.then, "enter") if edge == "enter" else []
    if trigger == "exit":
        return _actions_for_trigger(rule.then, "exit") if edge == "exit" else []
    if trigger == "change":
        return _actions_for_trigger(rule.then, edge) if edge in {"enter", "exit"} else []
    return []


def _actions_for_trigger(
    actions: list[ActionSpec] | ActionSet,
    trigger: Literal["while", "enter", "exit", "change"],
) -> list[ActionSpec]:
    if isinstance(actions, list):
        return actions if trigger in {"while", "enter", "exit"} else []
    if trigger == "while":
        return actions.while_actions or actions.enter
    if trigger == "enter":
        return actions.enter
    if trigger == "exit":
        return actions.exit
    return [*actions.enter, *actions.exit]