agents/message_protocol.py

from __future__ import annotations

import json
from dataclasses import asdict, dataclass, field
from typing import Any


VALID_MODES = {
    "none",
    "prioritize_ns",
    "prioritize_ew",
    "green_wave",
    "emergency_route",
    "damp_border_inflow",
}


@dataclass
class NeighborMessage:
    sender_intersection: str
    receiver_intersection: str
    congestion_level: float
    spillback_risk: bool
    dominant_direction: str  # "ns", "ew", or "balanced"
    queue_total: int

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)


@dataclass
class DistrictDirective:
    mode: str = "none"
    target_intersections: list[str] = field(default_factory=list)
    duration: int = 1
    rationale: str = ""
    corridor: str | None = None
    district_weight: float = 0.5

    def validate(self) -> "DistrictDirective":
        if self.mode not in VALID_MODES:
            self.mode = "none"

        if not isinstance(self.target_intersections, list):
            self.target_intersections = []

        if not isinstance(self.duration, int):
            self.duration = 1
        self.duration = max(1, min(self.duration, 10))

        if not isinstance(self.rationale, str):
            self.rationale = ""

        if self.corridor is not None and self.corridor not in {
            "ns",
            "ew",
            "west_to_east",
            "east_to_west",
            "north_to_south",
            "south_to_north",
        }:
            self.corridor = None

        if not isinstance(self.district_weight, (int, float)):
            self.district_weight = 0.5
        self.district_weight = float(max(0.0, min(1.0, self.district_weight)))

        return self

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)


def parse_district_directive(payload: str | dict[str, Any]) -> DistrictDirective:
    """
    Accept either raw JSON text or a dict and return a validated DistrictDirective.
    Falls back safely to a no-op directive.
    """
    try:
        if isinstance(payload, str):
            payload = payload.strip()
            if not payload:
                return DistrictDirective().validate()

            # Try direct JSON parse
            try:
                data = json.loads(payload)
            except json.JSONDecodeError:
                # Try to extract JSON object from surrounding text
                start = payload.find("{")
                end = payload.rfind("}")
                if start == -1 or end == -1 or end <= start:
                    return DistrictDirective().validate()
                data = json.loads(payload[start : end + 1])
        elif isinstance(payload, dict):
            data = payload
        else:
            return DistrictDirective().validate()

        directive = DistrictDirective(
            mode=data.get("mode", "none"),
            target_intersections=data.get("target_intersections", []),
            duration=data.get("duration", 1),
            rationale=data.get("rationale", ""),
            corridor=data.get("corridor"),
            district_weight=data.get("district_weight", 0.5),
        )
        return directive.validate()
    except Exception:
        return DistrictDirective().validate()


def safe_directive_dict(payload: str | dict[str, Any] | None) -> dict[str, Any]:
    if payload is None:
        return DistrictDirective().validate().to_dict()
    return parse_district_directive(payload).to_dict()