from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Tuple
from collections import deque

import numpy as np

from ..core.base import BaseProcessor, ProcessingContext, ProcessingResult, ConfigProtocol
from ..core.config import BaseConfig, AlertConfig

from ..utils import (
    filter_by_confidence,
    apply_category_mapping,
    match_results_structure,
    bbox_smoothing,
    BBoxSmoothingConfig,
    BBoxSmoothingTracker,
)

from ..utils.bytetrack_utils import (
    SORTTracker,
    ByteTrackWrapper,
    bbox_centroid,
    bbox_feet_point,
    dist,
    smooth_point,
    bbox_iou,
)


# =============================================================================
# Config
# =============================================================================
@dataclass
class LoiteringConfig(BaseConfig):
    """
    Loitering detection config (video-time aware).

    Loitering decision:
    - Person must be present >= loiter_threshold_seconds
    - Stationary ratio over sliding behavior window >= stationary_ratio_threshold
    - Avg speed in window <= velocity_threshold_px_per_sec

    Notes:
    - "Presence" is accumulated using video-time dt derived from fps (not wall-clock).
    - "Stationary ratio" is computed using a window of boolean flags based on instantaneous speed.
    """

    # --- Filtering ---
    confidence_threshold: float = 0.6
    target_categories: List[str] = field(default_factory=lambda: ["person"])

    # --- Loitering rules ---
    loiter_threshold_seconds: float = 10.0
    velocity_threshold_px_per_sec: float = 18.0
    stationary_ratio_threshold: float = 0.70
    min_presence_seconds: float = 2.0

    # --- Sliding behavior window (seconds) ---
    behavior_window_seconds: float = 8.0
    min_behavior_window_seconds: float = 4.0

    # --- Track lifecycle (video-time) ---
    track_timeout_seconds: float = 12.0
    missing_tolerance_seconds: float = 6.0

    # --- Re-association / resurrection ---
    resurrection_time_window_seconds: float = 10.0
    resurrection_centroid_dist_px: float = 180.0
    resurrection_iou_threshold: float = 0.10

    # --- Jitter robustness ---
    max_centroid_jump_px: float = 80.0
    centroid_ema_alpha: float = 0.25

    # --- Speed windows (frames) ---
    speed_window_size: int = 15
    slow_flags_window_size: int = 15

    # --- Optional bbox smoothing ---
    enable_smoothing: bool = True
    smoothing_algorithm: str = "observability"
    smoothing_window_size: int = 20
    smoothing_cooldown_frames: int = 5
    smoothing_confidence_range_factor: float = 0.5

    # --- Category mapping ---
    index_to_category: Optional[Dict[int, str]] = field(default_factory=lambda: {0: "person"})

    # --- Alerts ---
    alert_cooldown_seconds: float = 4.0
    alert_config: Optional[AlertConfig] = None

    # --- Tracking control ---
    enable_tracking: bool = True
    tracking_method: str = "sort"  # "sort" | "bytetrack"

    # Generic tracking knobs (used by SORT, and partially by ByteTrack)
    tracking_max_age: int = 30
    tracking_min_hits: int = 2
    tracking_iou_threshold: float = 0.25

    # ByteTrack knobs
    bytetrack_track_thresh: float = 0.25
    bytetrack_match_thresh: float = 0.80
    bytetrack_min_iou_assign: float = 0.30


# =============================================================================
# Use Case
# =============================================================================
class LoiteringUseCase(BaseProcessor):
    """
    Matrice Loitering UseCase (v3.1)

    Flow:
      1) Normalize input detections
      2) Filter + category mapping
      3) Optional bbox smoothing
      4) Tracking (SORT or ByteTrack)
      5) Update per-track loiter state using VIDEO time dt
      6) Mark detections with is_loitering and loitering_meta (ALWAYS)
      7) Create agg_summary + alerts

    Key outputs:
      - result.data["detections"] contains the tracked + labeled detections for visualization.
      - Every detection will have det["loitering_meta"]["loitering_confidence"] populated.
    """

    GLOBAL_ZONE_NAME = "global"

    def __init__(self):
        super().__init__("loitering")

        self.category = "general"
        self.CASE_TYPE: str = "loitering"
        self.CASE_VERSION: str = "3.1"

        # Categories to consider for loitering decision (typically just "person")
        self.target_categories = ["person"]

        # Optional bbox smoother (helps reduce jitter before tracking / speed estimation)
        self.smoothing_tracker: Optional[BBoxSmoothingTracker] = None

        # Optional internal tracker (SORT or YOLOX ByteTrack wrapper)
        # NOTE: In your YOLO.track runner, enable_tracking=False and track_id comes from YOLO.
        self.tracker: Optional[Any] = None

        # Internal frame counter (video-time is derived using fps + this counter)
        self._total_frame_counter = 0

        # Canonical per-track loiter state:
        #   track_id -> dict(state fields)
        self._loiter_tracks: Dict[int, Dict[str, Any]] = {}

        # Unique counting states (compatible with Matrice tracking_stats schema)
        self._per_category_total_track_ids: Dict[str, set] = {}
        self._current_frame_track_ids: Dict[str, set] = {}

    # -------------------------------------------------------------------------
    # VIDEO-time dt
    # -------------------------------------------------------------------------
    def _get_video_dt(self, stream_info: Optional[Dict[str, Any]]) -> float:
        """
        Compute dt (seconds per frame) based on the ORIGINAL input fps.
        This makes the logic "video-time aware" and deterministic across environments.
        """
        fps = 30.0
        try:
            if stream_info:
                fps_val = stream_info.get("input_settings", {}).get("original_fps")
                if fps_val and float(fps_val) > 1e-6:
                    fps = float(fps_val)
        except Exception:
            fps = 30.0

        return float(1.0 / max(1e-6, fps))

    # -------------------------------------------------------------------------
    # Tracking init
    # -------------------------------------------------------------------------
    def _init_tracker(self, config: LoiteringConfig, stream_info: Optional[Dict[str, Any]]) -> None:
        """
        Initialize an internal tracker if enabled by config.

        Supported:
          - SORTTracker (Kalman + Hungarian)
          - ByteTrackWrapper (YOLOX BYTETracker wrapper)

        NOTE:
        - If you are already running YOLO.track(persist=True), you typically set
          enable_tracking=False so this is skipped.
        """
        if self.tracker is not None:
            return

        method = str(getattr(config, "tracking_method", "sort")).lower().strip()

        if method == "sort":
            self.tracker = SORTTracker(
                iou_threshold=float(getattr(config, "tracking_iou_threshold", 0.25)),
                max_age=int(getattr(config, "tracking_max_age", 30)),
                min_hits=int(getattr(config, "tracking_min_hits", 2)),
            )
            return

        if method == "bytetrack":
            # ByteTrackWrapper needs fps mostly for track management / time scaling
            fps = 30.0
            try:
                if stream_info:
                    fps_val = stream_info.get("input_settings", {}).get("original_fps")
                    if fps_val and float(fps_val) > 1e-6:
                        fps = float(fps_val)
            except Exception:
                fps = 30.0

            self.tracker = ByteTrackWrapper(
                fps=float(fps),
                track_thresh=float(getattr(config, "bytetrack_track_thresh", 0.25)),
                match_thresh=float(getattr(config, "bytetrack_match_thresh", 0.80)),
                track_buffer=int(getattr(config, "tracking_max_age", 30)),
            )
            return

        # Unknown method => no tracking
        self.tracker = None

    # -------------------------------------------------------------------------
    # Schema normalization
    # -------------------------------------------------------------------------
    def _normalize_detection_schema(self, dets: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """
        Normalize incoming detection schema into a consistent shape.

        Accepted inputs may contain:
          - bbox or bounding_box (both supported)
          - category or class (both supported)
        """
        fixed: List[Dict[str, Any]] = []
        for d in dets:
            if not isinstance(d, dict):
                continue

            # unify bbox keys
            if "bounding_box" not in d and "bbox" in d:
                d["bounding_box"] = d["bbox"]
            if "bbox" not in d and "bounding_box" in d:
                d["bbox"] = d["bounding_box"]

            # unify category keys
            if "category" not in d and "class" in d:
                d["category"] = d["class"]
            if "class" not in d and "category" in d:
                d["class"] = d["category"]

            fixed.append(d)
        return fixed

    # -------------------------------------------------------------------------
    # Track state init/update
    # -------------------------------------------------------------------------
    def _init_track_state(
        self,
        bbox: Dict[str, Any],
        centroid: Tuple[float, float],
        feet: Tuple[float, float],
        config: LoiteringConfig,
    ) -> Dict[str, Any]:
        """
        Create a new per-track state dict.

        The two sliding windows:
          - speed_window: instantaneous feet speed per frame (px/sec)
          - slow_flags_window: 1.0 if stationary else 0.0 (same length as speed_window)
        """
        win = max(3, int(config.speed_window_size))
        return {
            "presence_seconds": 0.0,
            "missing_for_seconds": 0.0,
            "last_bbox": bbox,
            "last_centroid": centroid,
            "smoothed_centroid": centroid,
            "last_feet": feet,
            "smoothed_feet": feet,
            "speed_window": deque(maxlen=win),
            "slow_flags_window": deque(maxlen=max(3, int(config.slow_flags_window_size))),
            "last_inst_speed_feet": 0.0,
            "last_inst_speed_centroid": 0.0,
            "is_loitering": False,
            "last_alert_video_time": None,
            "resurrection_hits": 0,
        }

    def _is_stationary(self, inst_speed_feet: float, inst_speed_centroid: float, config: LoiteringConfig) -> bool:
        """
        Determine if a track is stationary for this frame.
        Feet speed is treated more strictly, centroid speed is allowed more tolerance.
        """
        vth = float(config.velocity_threshold_px_per_sec)
        return (inst_speed_feet <= vth * 0.80) and (inst_speed_centroid <= vth * 1.20)

    # -------------------------------------------------------------------------
    # Loiter confidence (0..1)
    # -------------------------------------------------------------------------
    def _compute_loiter_confidence(
        self,
        presence: float,
        win_slow_ratio: float,
        win_avg_speed: float,
        config: LoiteringConfig,
    ) -> float:
        """
        Compute loiter-confidence in [0..1] for EVERY person (even before loiter becomes true).

        Uses 3 normalized factors:
          - presence vs loiter_threshold
          - slow_ratio vs stationary_ratio_threshold
          - speed vs velocity_threshold (lower speed => higher confidence)

        This metric is meant for visualization/debugging and gradual buildup.
        """
        # 1) Presence factor: grows towards 1 when reaching loiter threshold
        p_score = min(1.0, presence / max(1e-6, float(config.loiter_threshold_seconds)))

        # 2) Stationary ratio factor: grows as slow ratio approaches the configured threshold
        sr_score = min(1.0, win_slow_ratio / max(1e-6, float(config.stationary_ratio_threshold)))

        # 3) Speed factor: 1 if speed <= threshold, else it decays
        v_th = float(config.velocity_threshold_px_per_sec)
        if win_avg_speed <= v_th:
            v_score = 1.0
        else:
            v_score = max(0.0, v_th / win_avg_speed)

        conf = float((p_score + sr_score + v_score) / 3.0)
        return max(0.0, min(1.0, conf))

    # -------------------------------------------------------------------------
    # Core state update (VIDEO TIME)
    # -------------------------------------------------------------------------
    def _update_loiter_states(
        self,
        detections: List[Dict[str, Any]],
        config: LoiteringConfig,
        dt_video: float,
        video_time_seconds: float,
    ) -> None:
        """
        Update the canonical per-track loiter state in `self._loiter_tracks`.

        This method is the "state machine heart" of the usecase.

        What it does:
        - Iterates detections for the configured target_categories ("person" by default)
        - For each detection with a valid track_id:
            * Initialize per-track state if this is a new track
            * Accumulate presence time using dt_video (fps-based)
            * Smooth centroid + feet point to reduce jitter (EMA)
            * Clamp large centroid jumps (prevents spikes in speed)
            * Compute instantaneous speed (px/sec) for centroid and feet
            * Push speed + stationary flags into sliding windows
            * Decide loitering based on:
                - presence >= loiter_threshold_seconds
                - slow_ratio >= stationary_ratio_threshold
                - avg_speed <= velocity_threshold_px_per_sec
        - Tracks missing detections and expires old tracks after track_timeout_seconds

        Notes:
        - `video_time_seconds` is currently used for alert cooldown outside this method,
          but it is kept in signature to support future time-based extensions.
        - This method does NOT modify detection objects directly. It only updates track state.
        """
        # Track IDs that appear in this frame (used to handle "missing" tracks later)
        present_ids: set[int] = set()

        # Helper: safely extract bbox dict from a detection
        def _get_bbox(det: Dict[str, Any]) -> Optional[Dict[str, float]]:
            bbox = det.get("bounding_box") or det.get("bbox")
            return bbox if isinstance(bbox, dict) else None

        # Helper: attempt to "heal" / merge an old track into a newly seen track_id
        # if the tracker re-assigns IDs briefly (common in occlusion / ID switches).
        def _try_heal_track_id(
            new_tid: int,
            new_bbox: Dict[str, float],
            new_feet: Tuple[float, float],
        ) -> Optional[int]:
            best_old_tid: Optional[int] = None
            best_score: float = -1.0

            # Only consider tracks missing for <= track_timeout_seconds
            max_missing = float(config.track_timeout_seconds)

            # Healing thresholds (soft-configurable via getattr)
            iou_thr = float(getattr(config, "id_heal_iou_threshold", 0.30))
            feet_dist_thr = float(getattr(config, "id_heal_feet_distance_px", 60.0))

            for old_tid, st in self._loiter_tracks.items():
                if old_tid == new_tid:
                    continue

                # Only consider "recently missing" tracks
                missing_for = float(st.get("missing_for_seconds", 0.0))
                if missing_for <= 0.0 or missing_for > max_missing:
                    continue

                old_bbox = st.get("last_bbox")
                if not isinstance(old_bbox, dict):
                    continue

                # Primary match signal = bbox IoU
                iou = float(bbox_iou(old_bbox, new_bbox))
                if iou >= iou_thr:
                    score = iou
                else:
                    # Fallback match signal = feet-point proximity
                    old_feet = st.get("smoothed_feet") or st.get("last_feet")
                    if not old_feet:
                        continue

                    feet_dist = float(dist(old_feet, new_feet))
                    if feet_dist <= feet_dist_thr:
                        score = 1.0 - min(1.0, feet_dist / max(1.0, feet_dist_thr))
                    else:
                        continue

                if score > best_score:
                    best_score = score
                    best_old_tid = old_tid

            if best_old_tid is None:
                return None

            # Merge old state into new track ID
            old_state = self._loiter_tracks.get(best_old_tid)
            if not old_state:
                return None

            self._loiter_tracks[new_tid] = old_state
            self._loiter_tracks.pop(best_old_tid, None)

            if getattr(config, "enable_id_healing_debug", True):
                self.logger.info(
                    f"[LOITER-ID-HEAL] merged old_tid={best_old_tid} -> new_tid={new_tid} score={best_score:.2f}"
                )

            return best_old_tid

        # ---------------------------------------------------------------------
        # Per-detection update (present tracks)
        # ---------------------------------------------------------------------
        for det in detections:
            # Ignore unrelated categories
            if det.get("category") not in self.target_categories:
                continue

            # Must have track_id for temporal analysis
            tid = det.get("track_id")
            if tid is None:
                continue

            try:
                tid_int = int(tid)
            except Exception:
                continue

            if tid_int < 0:
                continue

            bbox = _get_bbox(det)
            if not bbox:
                continue

            # Feature points derived from bbox
            centroid = bbox_centroid(bbox)
            feet = bbox_feet_point(bbox)

            # If track_id is new, attempt to heal by merging a recently-missing track
            if tid_int not in self._loiter_tracks and self._loiter_tracks:
                _try_heal_track_id(tid_int, bbox, feet)

            present_ids.add(tid_int)

            # If still new after healing attempt: initialize state and move on
            if tid_int not in self._loiter_tracks:
                self._loiter_tracks[tid_int] = self._init_track_state(bbox, centroid, feet, config)
                continue

            st = self._loiter_tracks[tid_int]

            # Presence accumulates in video-time (dt_video)
            st["presence_seconds"] = float(st.get("presence_seconds", 0.0) + dt_video)

            # Reset missing time since we saw it this frame
            st["missing_for_seconds"] = 0.0

            # Use smoothed points as speed reference (reduces jitter spikes)
            prev_centroid = st.get("smoothed_centroid", centroid)
            prev_feet = st.get("smoothed_feet", feet)

            # Exponential moving average smoothing for centroid
            centroid_alpha = float(config.centroid_ema_alpha)
            new_centroid = smooth_point(prev_centroid, centroid, centroid_alpha)

            # Feet point gets a slightly stronger smoothing by default
            feet_alpha = min(0.40, centroid_alpha + 0.10)
            new_feet = smooth_point(prev_feet, feet, float(feet_alpha))

            # Clamp sudden centroid jumps to keep speed stable
            max_jump = float(config.max_centroid_jump_px)
            if dist(prev_centroid, new_centroid) > max_jump:
                dx = new_centroid[0] - prev_centroid[0]
                dy = new_centroid[1] - prev_centroid[1]
                norm = max(1e-6, (dx * dx + dy * dy) ** 0.5)
                scale = max_jump / norm
                new_centroid = (prev_centroid[0] + dx * scale, prev_centroid[1] + dy * scale)

            # Instantaneous speeds (px/sec)
            dt_safe = max(1e-6, dt_video)
            inst_speed_centroid = dist(new_centroid, prev_centroid) / dt_safe
            inst_speed_feet = dist(new_feet, prev_feet) / dt_safe

            st["last_inst_speed_centroid"] = float(inst_speed_centroid)
            st["last_inst_speed_feet"] = float(inst_speed_feet)

            # Convert speeds into a stationary flag (1.0 or 0.0)
            stationary = self._is_stationary(inst_speed_feet, inst_speed_centroid, config)
            slow_flag = 1.0 if stationary else 0.0

            # Update sliding windows
            st["speed_window"].append(float(inst_speed_feet))
            st["slow_flags_window"].append(float(slow_flag))

            # Update last + smoothed anchors
            st["last_bbox"] = bbox
            st["last_centroid"] = centroid
            st["smoothed_centroid"] = new_centroid
            st["last_feet"] = feet
            st["smoothed_feet"] = new_feet

            presence = float(st.get("presence_seconds", 0.0))

            # Do not allow loitering decisions too early (warm-up)
            if presence < float(config.min_presence_seconds):
                st["is_loitering"] = False
                self._loiter_tracks[tid_int] = st
                continue

            # A minimal window is required for stable behavior stats
            enough_window = (
                len(st["speed_window"]) >= 3
                and len(st["slow_flags_window"]) >= 3
                and min(presence, float(config.behavior_window_seconds)) >= float(config.min_behavior_window_seconds)
            )

            # Compute behavior-window statistics
            win_avg_speed = float(np.mean(list(st["speed_window"]))) if st["speed_window"] else 0.0
            win_slow_ratio = float(np.mean(list(st["slow_flags_window"]))) if st["slow_flags_window"] else 0.0

            # Final loiter decision
            is_loitering = False
            if enough_window:
                is_loitering = (
                    presence >= float(config.loiter_threshold_seconds)
                    and win_slow_ratio >= float(config.stationary_ratio_threshold)
                    and win_avg_speed <= float(config.velocity_threshold_px_per_sec)
                )

            # Periodic debug (only after threshold time)
            if presence >= float(config.loiter_threshold_seconds) and (self._total_frame_counter % 50 == 0):
                self.logger.info(
                    f"[LOITER-DEBUG] tid={tid_int} presence={presence:.2f}s "
                    f"avg_speed={win_avg_speed:.2f} slow_ratio={win_slow_ratio:.2f} loiter={is_loitering}"
                )

            st["is_loitering"] = bool(is_loitering)
            self._loiter_tracks[tid_int] = st

        # ---------------------------------------------------------------------
        # Missing track handling (tracks not observed this frame)
        # ---------------------------------------------------------------------
        for tid, st in list(self._loiter_tracks.items()):
            if tid in present_ids:
                continue

            # Accumulate missing time in video-time
            st["missing_for_seconds"] = float(st.get("missing_for_seconds", 0.0) + dt_video)

            # Drop tracks that have been missing too long
            if float(st["missing_for_seconds"]) > float(config.track_timeout_seconds):
                self._loiter_tracks.pop(tid, None)
            else:
                self._loiter_tracks[tid] = st

    # -------------------------------------------------------------------------
    # Alerts
    # -------------------------------------------------------------------------
    def _check_alerts(
        self,
        detections: List[Dict[str, Any]],
        frame_key: str,
        config: LoiteringConfig,
        video_time_seconds: float,
    ) -> List[Dict[str, Any]]:
        """
        Generate loitering alerts (cooldown-enforced per track_id).

        Alerts are only emitted when:
          - a track is currently loitering
          - cooldown time has passed since last alert for the same track
        """
        alerts: List[Dict[str, Any]] = []

        # Default alert output format
        alert_type = ["Default"]
        alert_value = ["JSON"]
        settings_map = {"Default": "JSON"}

        # If alert_config exists, build type->value mapping
        if config.alert_config:
            alert_type = getattr(config.alert_config, "alert_type", ["Default"])
            alert_value = getattr(config.alert_config, "alert_value", ["JSON"])
            settings_map = {t: v for t, v in zip(alert_type, alert_value)}

        for det in detections:
            tid = det.get("track_id")
            if tid is None or int(tid) < 0:
                continue

            tid = int(tid)
            st = self._loiter_tracks.get(tid)
            if not st:
                continue

            # Only alert on loitering tracks
            if not bool(st.get("is_loitering", False)):
                continue

            # Cooldown check (video-time based)
            last_alert = st.get("last_alert_video_time", None)
            if last_alert is not None and (video_time_seconds - float(last_alert)) < float(config.alert_cooldown_seconds):
                continue

            # Mark alert time
            st["last_alert_video_time"] = float(video_time_seconds)
            self._loiter_tracks[tid] = st

            bbox = det.get("bounding_box") or det.get("bbox")

            speed_window = list(st.get("speed_window", []))
            slow_flags = list(st.get("slow_flags_window", []))
            win_avg_speed = float(np.mean(speed_window)) if speed_window else 0.0
            win_slow_ratio = float(np.mean(slow_flags)) if slow_flags else 0.0

            alerts.append(
                {
                    "alert_type": alert_type,
                    "alert_id": f"loitering_alert_{tid}_{frame_key}",
                    "incident_category": self.CASE_TYPE,
                    "track_id": tid,
                    "zone_name": self.GLOBAL_ZONE_NAME,
                    "bounding_box": bbox,
                    "confidence": float(det.get("confidence", 0.0)),
                    "category": det.get("category"),
                    "dwell_seconds": round(float(st.get("presence_seconds", 0.0)), 2),
                    "window_slow_ratio": round(float(win_slow_ratio), 3),
                    "avg_speed_px_per_sec": round(float(win_avg_speed), 3),
                    "threshold_seconds": float(config.loiter_threshold_seconds),
                    "settings": settings_map,
                }
            )

        return alerts

    # -------------------------------------------------------------------------
    # Counting
    # -------------------------------------------------------------------------
    def _update_tracking_state(self, detections: List[Dict[str, Any]]) -> None:
        """
        Update unique-counting state used in tracking_stats output.

        We maintain:
          - _per_category_total_track_ids: unique ids seen over time
          - _current_frame_track_ids: ids seen in this frame
        """
        categories = self.target_categories + ["loitering_person"]
        if not self._per_category_total_track_ids:
            self._per_category_total_track_ids = {cat: set() for cat in categories}

        self._current_frame_track_ids = {cat: set() for cat in categories}

        for det in detections:
            cat = det.get("category")
            tid = det.get("track_id")
            if cat not in categories or tid is None or int(tid) < 0:
                continue

            tid = int(tid)
            self._per_category_total_track_ids.setdefault(cat, set()).add(tid)
            self._current_frame_track_ids[cat].add(tid)

    def get_total_counts(self) -> Dict[str, int]:
        """Return total unique track_id counts per category."""
        return {cat: len(ids) for cat, ids in self._per_category_total_track_ids.items()}

    def _count_categories(self, detections: List[Dict[str, Any]]) -> Dict[str, Any]:
        """
        Frame-level raw detection count (not unique-tracks).

        This supports the Matrice agg_summary schema:
          - total_count
          - per_category_count
          - list of detections (minimal schema)
        """
        counts: Dict[str, int] = {}
        for det in detections:
            cat = det.get("category", "unknown")
            counts[cat] = counts.get(cat, 0) + 1

        return {
            "total_count": int(sum(counts.values())),
            "per_category_count": counts,
            "detections": [
                {
                    "bounding_box": det.get("bounding_box"),
                    "category": det.get("category"),
                    "confidence": det.get("confidence"),
                    "track_id": det.get("track_id"),
                    "frame_id": det.get("frame_id"),
                    "is_loitering": bool(det.get("is_loitering", False)),
                }
                for det in detections
            ],
        }

    # -------------------------------------------------------------------------
    # Main Process
    # -------------------------------------------------------------------------
    def process(
        self,
        data: Any,
        config: ConfigProtocol,
        context: Optional[ProcessingContext] = None,
        stream_info: Optional[Dict[str, Any]] = None,
    ) -> ProcessingResult:
        """
        Primary entrypoint.

        Input:
          - `data`: detection list (or nested dict of detections)
          - `config`: LoiteringConfig
          - `stream_info`: must contain original_fps for accurate dt

        Output:
          ProcessingResult with:
            data["agg_summary"] : Matrice-standard agg_summary payload
            data["detections"]  : list of tracked + labeled detections (for visualization)
        """
        # Validate config type
        if not isinstance(config, LoiteringConfig):
            return self.create_error_result(
                "Invalid config type",
                usecase=self.name,
                category=self.category,
                context=context,
            )

        # Ensure we always have a context object
        if context is None:
            context = ProcessingContext()

        # Detect input format for compatibility with other postprocessing modules
        if isinstance(data, list):
            context.input_format = "detection"
        else:
            try:
                context.input_format = match_results_structure(data)
            except Exception:
                context.input_format = "detection"

        # Carry config confidence into context for standardization
        context.confidence_threshold = float(config.confidence_threshold)

        # Extract frame number if available (used as frame_id key in agg_summary)
        frame_number: Optional[int] = None
        if stream_info:
            input_settings = stream_info.get("input_settings", {})
            start_frame = input_settings.get("start_frame")
            end_frame = input_settings.get("end_frame")
            if start_frame is not None and end_frame is not None and start_frame == end_frame:
                try:
                    frame_number = int(start_frame)
                except Exception:
                    frame_number = None

        frame_id = str(frame_number) if frame_number is not None else "current_frame"

        # Video-time scaling (fps-based)
        dt_video = self._get_video_dt(stream_info)
        video_time_seconds = float((self._total_frame_counter + 1) * dt_video)

        # ---------------------------------------------------------------------
        # 1) Parse raw detections from input
        # ---------------------------------------------------------------------
        detections: List[Dict[str, Any]] = []
        if isinstance(data, dict):
            # Nested schema: {"frame_id": [detections...], ...}
            for _, dets in data.items():
                if isinstance(dets, list):
                    detections.extend([d for d in dets if isinstance(d, dict)])
        elif isinstance(data, list):
            # Flat schema: [detections...]
            detections = [d for d in data if isinstance(d, dict)]
        else:
            detections = []

        # Normalize key names and schema
        detections = self._normalize_detection_schema(detections)

        # ---------------------------------------------------------------------
        # 2) Filter detections and map categories (class index -> label)
        # ---------------------------------------------------------------------
        detections = filter_by_confidence(detections, float(config.confidence_threshold))
        if getattr(config, "index_to_category", None):
            detections = apply_category_mapping(detections, config.index_to_category)

        detections = self._normalize_detection_schema(detections)

        # Apply target category filter (only keep "person" detections typically)
        self.target_categories = list(getattr(config, "target_categories", ["person"])) or ["person"]
        detections = [d for d in detections if d.get("category") in self.target_categories]

        # ---------------------------------------------------------------------
        # 3) Optional bbox smoothing (reduces jitter before speed estimation)
        # ---------------------------------------------------------------------
        if bool(getattr(config, "enable_smoothing", False)):
            if self.smoothing_tracker is None:
                smoothing_config = BBoxSmoothingConfig(
                    smoothing_algorithm=str(getattr(config, "smoothing_algorithm", "observability")),
                    window_size=int(getattr(config, "smoothing_window_size", 20)),
                    cooldown_frames=int(getattr(config, "smoothing_cooldown_frames", 5)),
                    confidence_threshold=float(getattr(config, "confidence_threshold", 0.6)),
                    confidence_range_factor=float(getattr(config, "smoothing_confidence_range_factor", 0.5)),
                    enable_smoothing=True,
                )
                self.smoothing_tracker = BBoxSmoothingTracker(smoothing_config)

            detections = bbox_smoothing(detections, self.smoothing_tracker.config, self.smoothing_tracker)

        detections = self._normalize_detection_schema(detections)

        # ---------------------------------------------------------------------
        # 4) Optional internal tracking (SORT / YOLOX ByteTrack wrapper)
        # ---------------------------------------------------------------------
        if bool(getattr(config, "enable_tracking", False)):
            try:
                self._init_tracker(config, stream_info)
                if self.tracker is not None and hasattr(self.tracker, "update"):
                    if isinstance(self.tracker, ByteTrackWrapper):
                        detections = self.tracker.update(detections, stream_info=stream_info)
                    else:
                        detections = self.tracker.update(detections)
            except Exception as e:
                self.logger.warning(f"Tracking failed: {e}")

        detections = self._normalize_detection_schema(detections)

        # ---------------------------------------------------------------------
        # 5) Update per-track loiter states (presence, speed windows, loiter flag)
        # ---------------------------------------------------------------------
        self._update_loiter_states(detections, config, dt_video=dt_video, video_time_seconds=video_time_seconds)

        # ---------------------------------------------------------------------
        # 6) Mark detections with loitering fields + ALWAYS attach meta/confidence
        # ---------------------------------------------------------------------
        for det in detections:
            tid = det.get("track_id")

            # No/invalid track -> still attach meta, but confidence stays 0
            if tid is None:
                det["is_loitering"] = False
                det["loitering_meta"] = {
                    "zone_name": self.GLOBAL_ZONE_NAME,
                    "presence_seconds": 0.0,
                    "window_slow_ratio": 0.0,
                    "window_avg_speed_px_per_sec": 0.0,
                    "behavior_window_seconds": float(config.behavior_window_seconds),
                    "loitering_confidence": 0.0,
                }
                continue

            try:
                tid = int(tid)
            except Exception:
                det["is_loitering"] = False
                det["loitering_meta"] = {
                    "zone_name": self.GLOBAL_ZONE_NAME,
                    "presence_seconds": 0.0,
                    "window_slow_ratio": 0.0,
                    "window_avg_speed_px_per_sec": 0.0,
                    "behavior_window_seconds": float(config.behavior_window_seconds),
                    "loitering_confidence": 0.0,
                }
                continue

            if tid < 0:
                det["is_loitering"] = False
                det["loitering_meta"] = {
                    "zone_name": self.GLOBAL_ZONE_NAME,
                    "presence_seconds": 0.0,
                    "window_slow_ratio": 0.0,
                    "window_avg_speed_px_per_sec": 0.0,
                    "behavior_window_seconds": float(config.behavior_window_seconds),
                    "loitering_confidence": 0.0,
                }
                continue

            # Pull track state and compute per-detection meta
            st = self._loiter_tracks.get(tid)

            presence = float(st.get("presence_seconds", 0.0)) if st else 0.0
            speed_window = list(st.get("speed_window", [])) if st else []
            slow_flags = list(st.get("slow_flags_window", [])) if st else []

            win_avg_speed = float(np.mean(speed_window)) if speed_window else 0.0
            win_slow_ratio = float(np.mean(slow_flags)) if slow_flags else 0.0

            loiter_conf = self._compute_loiter_confidence(
                presence=presence,
                win_slow_ratio=win_slow_ratio,
                win_avg_speed=win_avg_speed,
                config=config,
            )

            is_loiter = bool(st.get("is_loitering", False)) if st else False
            det["is_loitering"] = is_loiter

            # ✅ ALWAYS attach meta for every person
            det["loitering_meta"] = {
                "zone_name": self.GLOBAL_ZONE_NAME,
                "presence_seconds": round(presence, 2),
                "window_slow_ratio": round(win_slow_ratio, 3),
                "window_avg_speed_px_per_sec": round(win_avg_speed, 3),
                "behavior_window_seconds": float(config.behavior_window_seconds),
                "loitering_confidence": round(float(loiter_conf), 3),
            }

            # Category swap only when loiter becomes TRUE (helps visualization)
            if is_loiter:
                det["category"] = "loitering_person"

        # ---------------------------------------------------------------------
        # 7) Update counting + build agg_summary and alerts
        # ---------------------------------------------------------------------
        self._update_tracking_state(detections)
        self._total_frame_counter += 1

        counting_summary = self._count_categories(detections)
        counting_summary["total_counts"] = self.get_total_counts()

        alerts = self._check_alerts(detections, frame_id, config, video_time_seconds=video_time_seconds)

        camera_info = self.get_camera_info_from_stream(stream_info)
        current_ts = self._get_current_timestamp_str(stream_info)
        start_ts = current_ts

        loiterers_in_frame = int(counting_summary.get("per_category_count", {}).get("loitering_person", 0))

        # Human-readable status text (included in incidents/tracking_stats)
        human_text = (
            f"LOITERING @ {current_ts}\n"
            f"Loiterers in frame: {loiterers_in_frame}\n"
            f"Threshold: {float(config.loiter_threshold_seconds):.1f}s | "
            f"v_th: {float(config.velocity_threshold_px_per_sec):.1f}px/s | "
            f"ratio_th: {float(config.stationary_ratio_threshold):.2f}"
        )

        # Create incident only if loitering is detected in the current frame
        incidents: Dict[str, Any] = {}
        if loiterers_in_frame > 0:
            incidents = self.create_incident(
                incident_id=f"{self.CASE_TYPE}_{frame_id}",
                incident_type=self.CASE_TYPE,
                severity_level="medium",
                human_text=human_text,
                camera_info=camera_info,
                alerts=alerts,
                alert_settings=[],
                start_time=start_ts,
                end_time="Incident still active",
                level_settings={"low": 1, "medium": 3, "significant": 5, "critical": 8},
            )

        # Build detections array in Matrice standardized tracking_stats schema
        detections_for_stats = []
        for det in counting_summary.get("detections", []):
            bbox = det.get("bounding_box", {}) or {}
            category = det.get("category", "unknown") or "unknown"
            detections_for_stats.append(self.create_detection_object(category, bbox))

        tracking_stats = self.create_tracking_stats(
            total_counts=[{"category": k, "count": int(v)} for k, v in self.get_total_counts().items()],
            current_counts=[
                {"category": k, "count": int(v)}
                for k, v in counting_summary.get("per_category_count", {}).items()
            ],
            detections=detections_for_stats,
            human_text=human_text,
            camera_info=camera_info,
            alerts=alerts,
            alert_settings=[],
            reset_settings=[{"interval_type": "daily", "reset_time": {"value": 9, "time_unit": "hour"}}],
            start_time=current_ts,
            reset_time=start_ts,
        )

        agg_summary = {
            frame_id: {
                "incidents": incidents,
                "tracking_stats": tracking_stats,
                "business_analytics": {},
                "alerts": alerts,
                "zone_analysis": {},
                "human_text": human_text,
            }
        }

        # Mark processing completion for context performance stats
        context.mark_completed()

        # Final standardized output
        return self.create_result(
            data={
                "agg_summary": agg_summary,
                "detections": detections,  # ✅ tracked + labeled detections (for visualization)
            },
            usecase=self.name,
            category=self.category,
            context=context,
        )

    # -------------------------------------------------------------------------
    # Timestamp helpers
    # -------------------------------------------------------------------------
    def _format_timestamp(self, timestamp: Any) -> str:
        """
        Keep timestamps short/consistent.

        If a string timestamp has fractional seconds, we clamp it to 2 decimals.
        """
        if not isinstance(timestamp, str):
            return str(timestamp)

        if "." not in timestamp:
            return timestamp

        main_part, fractional_and_suffix = timestamp.split(".", 1)

        if " " in fractional_and_suffix:
            fractional_part, suffix = fractional_and_suffix.split(" ", 1)
            suffix = " " + suffix
        else:
            fractional_part, suffix = fractional_and_suffix, ""

        fractional_part = (fractional_part + "00")[:2]
        return f"{main_part}.{fractional_part}{suffix}"

    def _get_current_timestamp_str(self, stream_info: Optional[Dict[str, Any]]) -> str:
        """Extract stream_time from stream_info and format it."""
        if not stream_info:
            return "NA"
        return self._format_timestamp(stream_info.get("input_settings", {}).get("stream_time", "NA"))