from dataclasses import dataclass

from edge.tracking import Track
from edge.geometry import (
    torso_angle_deg, head_below_hips, core_visible, visible_count,
    track_displacement,
)


@dataclass
class FallConfig:
    falling_angle: float = 45.0
    fallen_angle: float = 70.0
    aspect_ratio: float = 1.25
    height_collapse: float = 0.35   # baseline 대비 남은 높이 비율(이 값 미만이면 급붕괴)
    prone_confirm_s: float = 3.0
    no_motion_s: float = 10.0
    motion_thr: float = 0.02        # 정규화 관절변위 임계
    min_conf: float = 0.5
    min_visible: int = 10
    recover_grace_s: float = 1.0    # 사고 상태에서 회복 판정까지 요구되는 지속 비-prone 시간


@dataclass
class _State:
    name: str = "STANDING"
    baseline_height: float | None = None
    prone_since: float | None = None
    still_since: float | None = None
    fall_emitted: bool = False
    no_motion_emitted: bool = False
    not_prone_since: float | None = None


class FallDetector:
    def __init__(self, cfg: FallConfig):
        self.cfg = cfg
        self._states: dict[int, _State] = {}

    def forget(self, track_id: int) -> None:
        self._states.pop(track_id, None)

    def _prone(self, kps, bbox, st: _State, kp_ok: bool) -> tuple[bool, dict]:
        ar = bbox.aspect
        remaining = 1.0
        if st.baseline_height:
            remaining = bbox.height / st.baseline_height
        # bbox cues need NO keypoints — a fallen person is wide-and-short, or their bbox height
        # collapsed. These stay live even when the pose model has lost confidence (occluded /
        # on the floor), which is exactly when a fall must not be missed.
        prone = ar > self.cfg.aspect_ratio or remaining < self.cfg.height_collapse
        angle = torso_angle_deg(kps)
        # keypoint cues only when the keypoints are actually trustworthy, else they are noise.
        if kp_ok:
            prone = prone or angle > self.cfg.fallen_angle or head_below_hips(kps)
        return prone, {"torso_angle": round(angle, 1),
                       "aspect_ratio": round(ar, 2),
                       "height_remaining": round(remaining, 2)}

    def update(self, track: Track, now: float) -> tuple[str | None, dict]:
        st = self._states.setdefault(track.track_id, _State())
        f = track.latest
        kps, bbox = f.keypoints, f.bbox

        # Confidence gate — no longer a blanket early-return (that disabled fall/no_motion/
        # ghost for the occluded person on the ground). It now only decides whether the
        # KEYPOINT-based cues are trustworthy; the bbox cues always run.
        kp_ok = core_visible(kps, self.cfg.min_conf) and visible_count(kps) >= self.cfg.min_visible

        prone, meta = self._prone(kps, bbox, st, kp_ok)

        # incident 비활성일 때만 baseline(정상 키높이) 갱신 = incident 중 freeze. Only trust a
        # confident standing frame for the baseline so a bad low-conf bbox can't corrupt it.
        if st.name == "STANDING" and not prone and kp_ok:
            st.baseline_height = bbox.height

        if not prone:
            if st.name in ("FALLING", "FALLEN", "NO_MOTION"):
                # 단일 튐(noisy) 프레임만으로 CRITICAL 사고를 즉시 취소하지 않도록
                # recover_grace_s만큼 지속된 비-prone을 요구한다.
                if st.not_prone_since is None:
                    st.not_prone_since = now
                    return None, {}
                if now - st.not_prone_since >= self.cfg.recover_grace_s:
                    self._states[track.track_id] = _State(baseline_height=bbox.height)
                    return "recover", {}
                return None, {}
            st.prone_since = None
            return None, {}

        # prone 상태 (재-prone 프레임은 대기 중인 recover 유예를 취소)
        st.not_prone_since = None
        if st.prone_since is None:
            st.prone_since = now
            st.name = "FALLING"
        if now - st.prone_since >= self.cfg.prone_confirm_s and not st.fall_emitted:
            st.name = "FALLEN"
            st.fall_emitted = True
            st.still_since = now
            return "fall", meta

        # FALLEN 이후 미동 감시
        if st.fall_emitted:
            if track_displacement(track) < self.cfg.motion_thr:
                if st.still_since is None:
                    st.still_since = now
                elif now - st.still_since >= self.cfg.no_motion_s and not st.no_motion_emitted:
                    st.name = "NO_MOTION"
                    st.no_motion_emitted = True
                    return "no_motion", {**meta, "immobility_s": round(now - st.still_since, 1)}
            else:
                st.still_since = None
        return None, {}