"""
Lateral Edge Guard uses the model's lateral road-edge geometry to withhold lane
changes that lack room for a target lane. The model standard deviation remains
in metres: measurements above the validity limit are rejected, while valid
measurements use a one-sigma lower confidence bound for conservative clearance.
Unavailable geometry briefly holds the last output, then fails open because a
model dropout is not geometric evidence of a nearby edge. A visible outer lane
line on the target side is direct evidence that a lane exists and overrides the
edge-distance inference.
"""
from __future__ import annotations

import math
from dataclasses import dataclass, replace
from enum import IntEnum
from typing import Any

from iqpilot.cereal import custom, log
from iqpilot.common.constants import CV
from iqpilot.common.swaglog import cloudlog


MIN_ACTIVE_SPEED_MPS = 20.0 * CV.MPH_TO_MS  # Matches the lane-change speed gate and excludes parking manoeuvres.
MAX_VALID_ROAD_EDGE_STD_M = 1.0  # A 2-sigma bound beyond 2 m cannot distinguish an adjacent 3.5 m lane reliably.
# roadEdgeStd describes a single edge point, but it is applied to a 5-40 m minimum that already absorbs the
# spatial worst case; 1 sigma covers ~1.1x the measured p99 frame-to-frame spread of that minimum, 2 sigma 2.2x.
EDGE_CONFIDENCE_SIGMA = 1.0
ROAD_EDGE_LOOKAHEAD_MIN_M = 5.0  # Ignore near-field edge points dominated by vehicle-body perspective.
ROAD_EDGE_LOOKAHEAD_MAX_M = 40.0  # Covers about 2 s at the 20 m/s model-training reference speed.
LANE_CENTER_OFFSET_M = 3.5  # Typical freeway lane width and the target-centre lateral displacement.
# CarParams exposes neither width nor track; 0.95 m is half of an assumed conservative 1.90 m body width.
VEHICLE_LATERAL_HALF_WIDTH_M = 1.90 / 2.0
EDGE_CLEARANCE_MARGIN_M = 0.25  # Additional lateral separation between the vehicle body and detected road edge.
ADJACENT_LANE_LINE_PROB = 0.5
EGO_LANE_LINE_PROB_MIN = 0.5
MIN_MEASURED_LANE_WIDTH_M = 2.5
MAX_MEASURED_LANE_WIDTH_M = 4.5
# modelV2 lane lines are ordered outer-left, ego-left, ego-right, outer-right.
OUTER_LANE_LINE_INDEX = (0, 3)
EGO_LANE_LINE_INDEX = (1, 2)
REQUIRED_ROAD_EDGE_DISTANCE_M = LANE_CENTER_OFFSET_M + VEHICLE_LATERAL_HALF_WIDTH_M + EDGE_CLEARANCE_MARGIN_M
BLOCK_DEBOUNCE_S = 0.30  # Six model frames reject a transient close-edge prediction before blocking.
CLEAR_DEBOUNCE_S = 0.50  # Ten model frames make release slower than assertion for conservative hysteresis.
UNAVAILABLE_HOLD_S = 0.50  # Ten model frames bridge a short model-data dropout before failing open.
TIMER_EPSILON_S = 1e-9  # Floating-point comparison tolerance, far below one model tick.

LaneChangeDirection = log.LaneChangeDirection
LateralEdgeBlock = custom.IQLateralEdgeBlock


class RoadEdgeDataState(IntEnum):
  VALID = 0
  UNAVAILABLE = 1
  INVALID = 2


@dataclass(frozen=True, slots=True)
class RoadEdgeMeasurement:
  state: RoadEdgeDataState
  lateral_distance_m: float | None = None
  conservative_distance_m: float | None = None
  should_block: bool | None = None


@dataclass(frozen=True, slots=True)
class _SideState:
  blocked: bool = False
  block_timer_s: float = 0.0
  clear_timer_s: float = 0.0
  unavailable_timer_s: float = 0.0
  fallback_reported: bool = False


def evaluate_road_edge(edge: Any, std_m: Any, direction: int,
                       lane_width_m: float = LANE_CENTER_OFFSET_M) -> RoadEdgeMeasurement:
  if edge is None or std_m is None:
    return RoadEdgeMeasurement(RoadEdgeDataState.UNAVAILABLE)

  try:
    xs = edge.x
    ys = edge.y
    count = len(xs)
    y_count = len(ys)
  except (AttributeError, TypeError):
    return RoadEdgeMeasurement(RoadEdgeDataState.UNAVAILABLE)

  if count == 0 or y_count != count:
    return RoadEdgeMeasurement(RoadEdgeDataState.UNAVAILABLE)

  try:
    std = float(std_m)
  except (TypeError, ValueError):
    return RoadEdgeMeasurement(RoadEdgeDataState.INVALID)
  if not math.isfinite(std) or std < 0.0 or std > MAX_VALID_ROAD_EDGE_STD_M:
    return RoadEdgeMeasurement(RoadEdgeDataState.INVALID)

  lateral_distance_m: float | None = None
  for idx in range(count):
    try:
      x_m = float(xs[idx])
      y_m = float(ys[idx])
    except (IndexError, TypeError, ValueError):
      return RoadEdgeMeasurement(RoadEdgeDataState.UNAVAILABLE)
    if not math.isfinite(x_m) or not math.isfinite(y_m):
      return RoadEdgeMeasurement(RoadEdgeDataState.UNAVAILABLE)
    if not ROAD_EDGE_LOOKAHEAD_MIN_M <= x_m <= ROAD_EDGE_LOOKAHEAD_MAX_M:
      continue
    if ((direction == LaneChangeDirection.left and y_m >= 0.0) or
        (direction == LaneChangeDirection.right and y_m <= 0.0)):
      return RoadEdgeMeasurement(RoadEdgeDataState.INVALID)
    distance_m = abs(y_m)
    lateral_distance_m = distance_m if lateral_distance_m is None else min(lateral_distance_m, distance_m)

  if lateral_distance_m is None:
    return RoadEdgeMeasurement(RoadEdgeDataState.UNAVAILABLE)

  conservative_distance_m = lateral_distance_m - EDGE_CONFIDENCE_SIGMA * std
  required_distance_m = lane_width_m + VEHICLE_LATERAL_HALF_WIDTH_M + EDGE_CLEARANCE_MARGIN_M
  return RoadEdgeMeasurement(
    RoadEdgeDataState.VALID,
    lateral_distance_m,
    conservative_distance_m,
    conservative_distance_m < required_distance_m,
  )


def step_side_guard(state: _SideState, measurement: RoadEdgeMeasurement, speed_active: bool,
                    dt_s: float) -> tuple[_SideState, bool]:
  if not speed_active:
    return _SideState(), False

  if measurement.state == RoadEdgeDataState.UNAVAILABLE:
    unavailable_timer_s = state.unavailable_timer_s + dt_s
    if unavailable_timer_s < UNAVAILABLE_HOLD_S - TIMER_EPSILON_S:
      return _SideState(state.blocked, unavailable_timer_s=unavailable_timer_s,
                        fallback_reported=state.fallback_reported), False
    fallback_started = not state.fallback_reported
    return _SideState(unavailable_timer_s=unavailable_timer_s, fallback_reported=True), fallback_started

  should_block = bool(measurement.should_block) if measurement.state == RoadEdgeDataState.VALID else False
  if should_block == state.blocked:
    return _SideState(blocked=state.blocked), False

  if should_block:
    block_timer_s = state.block_timer_s + dt_s
    if block_timer_s >= BLOCK_DEBOUNCE_S - TIMER_EPSILON_S:
      return _SideState(blocked=True), False
    return _SideState(block_timer_s=block_timer_s), False

  clear_timer_s = state.clear_timer_s + dt_s
  if clear_timer_s >= CLEAR_DEBOUNCE_S - TIMER_EPSILON_S:
    return _SideState(), False
  return _SideState(blocked=True, clear_timer_s=clear_timer_s), False


class LateralEdgeGuard:
  def __init__(self) -> None:
    self._left = _SideState()
    self._right = _SideState()
    self.left_measurement = RoadEdgeMeasurement(RoadEdgeDataState.UNAVAILABLE)
    self.right_measurement = RoadEdgeMeasurement(RoadEdgeDataState.UNAVAILABLE)

  @staticmethod
  def _model_side(modeldata: Any, side_index: int) -> tuple[Any | None, Any | None]:
    if modeldata is None:
      return None, None
    try:
      edges = modeldata.roadEdges
      stds = modeldata.roadEdgeStds
      if len(edges) <= side_index or len(stds) <= side_index:
        return None, None
      return edges[side_index], stds[side_index]
    except (AttributeError, TypeError):
      return None, None

  @staticmethod
  def _lane_line_prob(modeldata: Any, index: int) -> float | None:
    if modeldata is None:
      return None
    try:
      probs = modeldata.laneLineProbs
      if len(probs) <= index:
        return None
      value = float(probs[index])
    except (AttributeError, TypeError, IndexError, ValueError):
      return None
    return value if math.isfinite(value) else None

  @classmethod
  def _adjacent_lane_visible(cls, modeldata: Any, side_index: int) -> bool:
    prob = cls._lane_line_prob(modeldata, OUTER_LANE_LINE_INDEX[side_index])
    return prob is not None and prob > ADJACENT_LANE_LINE_PROB

  @classmethod
  def _measured_lane_width(cls, modeldata: Any) -> float:
    left_prob = cls._lane_line_prob(modeldata, EGO_LANE_LINE_INDEX[0])
    right_prob = cls._lane_line_prob(modeldata, EGO_LANE_LINE_INDEX[1])
    if left_prob is None or right_prob is None:
      return LANE_CENTER_OFFSET_M
    if left_prob <= EGO_LANE_LINE_PROB_MIN or right_prob <= EGO_LANE_LINE_PROB_MIN:
      return LANE_CENTER_OFFSET_M
    try:
      lines = modeldata.laneLines
      left_y = float(lines[EGO_LANE_LINE_INDEX[0]].y[0])
      right_y = float(lines[EGO_LANE_LINE_INDEX[1]].y[0])
    except (AttributeError, TypeError, IndexError, ValueError):
      return LANE_CENTER_OFFSET_M
    width = abs(right_y - left_y)
    if not math.isfinite(width):
      return LANE_CENTER_OFFSET_M
    return min(max(width, MIN_MEASURED_LANE_WIDTH_M), MAX_MEASURED_LANE_WIDTH_M)

  @staticmethod
  def _apply_lane_evidence(measurement: RoadEdgeMeasurement, lane_visible: bool) -> RoadEdgeMeasurement:
    if lane_visible and measurement.state == RoadEdgeDataState.VALID and measurement.should_block:
      return replace(measurement, should_block=False)
    return measurement

  def update(self, modeldata: Any, v_ego_mps: float, dt_s: float) -> None:
    dt = max(float(dt_s), 0.0)
    left_edge, left_std = self._model_side(modeldata, 0)
    right_edge, right_std = self._model_side(modeldata, 1)
    lane_width_m = self._measured_lane_width(modeldata)
    self.left_measurement = self._apply_lane_evidence(
      evaluate_road_edge(left_edge, left_std, LaneChangeDirection.left, lane_width_m),
      self._adjacent_lane_visible(modeldata, 0))
    self.right_measurement = self._apply_lane_evidence(
      evaluate_road_edge(right_edge, right_std, LaneChangeDirection.right, lane_width_m),
      self._adjacent_lane_visible(modeldata, 1))
    speed_active = math.isfinite(v_ego_mps) and v_ego_mps >= MIN_ACTIVE_SPEED_MPS
    self._left, left_fallback = step_side_guard(self._left, self.left_measurement, speed_active, dt)
    self._right, right_fallback = step_side_guard(self._right, self.right_measurement, speed_active, dt)
    if left_fallback:
      cloudlog.warning(f"lateral edge guard: left road edge unavailable for {UNAVAILABLE_HOLD_S:.2f} s; falling back to not blocking")
    if right_fallback:
      cloudlog.warning(f"lateral edge guard: right road edge unavailable for {UNAVAILABLE_HOLD_S:.2f} s; falling back to not blocking")

  def block_for_direction(self, direction: int) -> custom.IQLateralEdgeBlock:
    if direction == LaneChangeDirection.left and self._left.blocked:
      return LateralEdgeBlock.left
    if direction == LaneChangeDirection.right and self._right.blocked:
      return LateralEdgeBlock.right
    return LateralEdgeBlock.none