#include "action_recog_node.h"

#include <algorithm>
#include <cmath>
#include <cstdint>
#include <deque>
#include <limits>
#include <map>
#include <memory>
#include <set>
#include <string>
#include <utility>
#include <vector>

#include "behavior/behavior_event.h"
#include "pose/pose_result.h"
#include "utils/logger.h"

namespace rk3588 {
namespace {

enum class ActionEventKind {
    Fall,
    Fight
};

enum class ActionSignalMode {
    Any,
    All
};

struct FallBboxRule {
    bool enabled = true;
    float min_drop_pixels = 0.0f;
    float min_aspect_ratio_delta = 0.0f;
};

struct FallPoseRule {
    bool enabled = false;
    float min_torso_drop_pixels = 0.0f;
    float max_upright_ratio = 0.0f;
};

struct FightBboxRule {
    bool enabled = true;
    float proximity_pixels = 0.0f;
    float min_motion_pixels = 0.0f;
};

struct FightPoseRule {
    bool enabled = false;
    float min_wrist_motion_pixels = 0.0f;
    float max_wrist_distance_pixels = 0.0f;
};

struct ActionEventRule {
    ActionEventKind kind = ActionEventKind::Fall;
    uint64_t window_ms = 1000;
    uint64_t activate_duration_ms = 0;
    ActionSignalMode signal_mode = ActionSignalMode::Any;
    FallBboxRule fall_bbox{};
    FallPoseRule fall_pose{};
    FightBboxRule fight_bbox{};
    FightPoseRule fight_pose{};
};

struct PoseTrackSample {
    bool matched = false;
    float torso_center_y = 0.0f;
    float upright_ratio = 0.0f;
    float left_wrist_x = 0.0f;
    float left_wrist_y = 0.0f;
    float right_wrist_x = 0.0f;
    float right_wrist_y = 0.0f;
    bool has_left_wrist = false;
    bool has_right_wrist = false;
};

struct TrackSample {
    uint64_t pts = 0;
    Rect bbox{};
    PoseTrackSample pose{};
};

static ActionSignalMode ParseSignalMode(const std::string& value) {
    if (value == "all") return ActionSignalMode::All;
    return ActionSignalMode::Any;
}

static bool EvaluateSignalMode(ActionSignalMode mode, bool bbox_trigger, bool pose_trigger,
                               bool bbox_enabled, bool pose_enabled) {
    if (bbox_enabled && pose_enabled) {
        return mode == ActionSignalMode::All ? (bbox_trigger && pose_trigger) : (bbox_trigger || pose_trigger);
    }
    if (bbox_enabled) return bbox_trigger;
    if (pose_enabled) return pose_trigger;
    return false;
}

static bool ParseFallRule(const SimpleJson& ev, ActionEventRule& rule) {
    rule.kind = ActionEventKind::Fall;
    rule.fall_bbox.min_drop_pixels = ev.ValueOr<float>("min_drop_pixels", 0.0f);
    rule.fall_bbox.min_aspect_ratio_delta = ev.ValueOr<float>("min_aspect_ratio_delta", 0.0f);
    rule.fall_pose.min_torso_drop_pixels = ev.ValueOr<float>("pose_min_torso_drop_pixels", 0.0f);
    rule.fall_pose.max_upright_ratio = ev.ValueOr<float>("pose_max_upright_ratio", 0.0f);

    if (const SimpleJson* bbox = ev.Find("bbox"); bbox && bbox->IsObject()) {
        rule.fall_bbox.enabled = bbox->ValueOr<bool>("enabled", true);
        rule.fall_bbox.min_drop_pixels = bbox->ValueOr<float>("min_drop_pixels", rule.fall_bbox.min_drop_pixels);
        rule.fall_bbox.min_aspect_ratio_delta =
            bbox->ValueOr<float>("min_aspect_ratio_delta", rule.fall_bbox.min_aspect_ratio_delta);
    } else {
        rule.fall_bbox.enabled = true;
    }

    if (const SimpleJson* pose = ev.Find("pose"); pose && pose->IsObject()) {
        rule.fall_pose.enabled = pose->ValueOr<bool>("enabled", true);
        rule.fall_pose.min_torso_drop_pixels =
            pose->ValueOr<float>("min_torso_drop_pixels", rule.fall_pose.min_torso_drop_pixels);
        rule.fall_pose.max_upright_ratio =
            pose->ValueOr<float>("max_upright_ratio", rule.fall_pose.max_upright_ratio);
    } else {
        rule.fall_pose.enabled = rule.fall_pose.min_torso_drop_pixels > 0.0f || rule.fall_pose.max_upright_ratio > 0.0f;
    }
    return true;
}

static bool ParseFightRule(const SimpleJson& ev, ActionEventRule& rule) {
    rule.kind = ActionEventKind::Fight;
    rule.fight_bbox.proximity_pixels = ev.ValueOr<float>("proximity_pixels", 0.0f);
    rule.fight_bbox.min_motion_pixels = ev.ValueOr<float>("min_motion_pixels", 0.0f);
    rule.fight_pose.min_wrist_motion_pixels = ev.ValueOr<float>("pose_min_wrist_motion_pixels", 0.0f);
    rule.fight_pose.max_wrist_distance_pixels = ev.ValueOr<float>("pose_max_wrist_distance_pixels", 0.0f);

    if (const SimpleJson* bbox = ev.Find("bbox"); bbox && bbox->IsObject()) {
        rule.fight_bbox.enabled = bbox->ValueOr<bool>("enabled", true);
        rule.fight_bbox.proximity_pixels = bbox->ValueOr<float>("proximity_pixels", rule.fight_bbox.proximity_pixels);
        rule.fight_bbox.min_motion_pixels = bbox->ValueOr<float>("min_motion_pixels", rule.fight_bbox.min_motion_pixels);
    } else {
        rule.fight_bbox.enabled = true;
    }

    if (const SimpleJson* pose = ev.Find("pose"); pose && pose->IsObject()) {
        rule.fight_pose.enabled = pose->ValueOr<bool>("enabled", true);
        rule.fight_pose.min_wrist_motion_pixels =
            pose->ValueOr<float>("min_wrist_motion_pixels", rule.fight_pose.min_wrist_motion_pixels);
        rule.fight_pose.max_wrist_distance_pixels =
            pose->ValueOr<float>("max_wrist_distance_pixels", rule.fight_pose.max_wrist_distance_pixels);
    } else {
        rule.fight_pose.enabled =
            rule.fight_pose.min_wrist_motion_pixels > 0.0f || rule.fight_pose.max_wrist_distance_pixels > 0.0f;
    }
    return true;
}

static bool ParseRules(const SimpleJson& config, std::vector<ActionEventRule>& out, std::string& err) {
    const SimpleJson* events = config.Find("events");
    if (!events || !events->IsArray()) {
        err = "events must be array";
        return false;
    }

    out.clear();
    for (const auto& ev : events->AsArray()) {
        if (!ev.IsObject()) {
            err = "event entry must be object";
            return false;
        }

        ActionEventRule rule;
        const std::string type = ev.ValueOr<std::string>("type", "");
        rule.window_ms = static_cast<uint64_t>(std::max(0, ev.ValueOr<int>("window_ms", 1000)));
        rule.activate_duration_ms = static_cast<uint64_t>(std::max(0, ev.ValueOr<int>("activate_duration_ms", 0)));
        rule.signal_mode = ParseSignalMode(ev.ValueOr<std::string>("signal_mode", "any"));
        if (const SimpleJson* fusion = ev.Find("fusion"); fusion && fusion->IsObject()) {
            rule.signal_mode = ParseSignalMode(fusion->ValueOr<std::string>("match_mode", "any"));
        }

        if (type == "fall") {
            ParseFallRule(ev, rule);
        } else if (type == "fight") {
            ParseFightRule(ev, rule);
        } else {
            err = "unsupported event type: " + type;
            return false;
        }

        out.push_back(std::move(rule));
    }
    return true;
}

static float CenterX(const Rect& rect) {
    return rect.x + (rect.w * 0.5f);
}

static float CenterY(const Rect& rect) {
    return rect.y + (rect.h * 0.5f);
}

static float CenterDistance(const Rect& lhs, const Rect& rhs) {
    const float dx = CenterX(lhs) - CenterX(rhs);
    const float dy = CenterY(lhs) - CenterY(rhs);
    return std::sqrt((dx * dx) + (dy * dy));
}

static float MotionDistance(const Rect& lhs, const Rect& rhs) {
    const float dx = CenterX(lhs) - CenterX(rhs);
    const float dy = CenterY(lhs) - CenterY(rhs);
    return std::sqrt((dx * dx) + (dy * dy));
}

static float AspectRatio(const Rect& rect) {
    return rect.h > 0.0f ? (rect.w / rect.h) : 0.0f;
}

static float BboxIoU(const Rect& lhs, const Rect& rhs) {
    const float x1 = std::max(lhs.x, rhs.x);
    const float y1 = std::max(lhs.y, rhs.y);
    const float x2 = std::min(lhs.x + lhs.w, rhs.x + rhs.w);
    const float y2 = std::min(lhs.y + lhs.h, rhs.y + rhs.h);
    const float iw = std::max(0.0f, x2 - x1);
    const float ih = std::max(0.0f, y2 - y1);
    const float inter = iw * ih;
    const float uni = lhs.w * lhs.h + rhs.w * rhs.h - inter;
    return uni <= 0.0f ? 0.0f : (inter / uni);
}

static bool IsValidKeypoint(const PoseItem& pose, int index) {
    return index >= 0 &&
           static_cast<size_t>(index) < pose.keypoints.size() &&
           pose.keypoints[static_cast<size_t>(index)].score > 0.0f;
}

static PoseTrackSample BuildPoseTrackSample(const PoseItem& pose) {
    PoseTrackSample sample;
    sample.matched = true;

    const bool has_left_shoulder = IsValidKeypoint(pose, 5);
    const bool has_right_shoulder = IsValidKeypoint(pose, 6);
    const bool has_left_hip = IsValidKeypoint(pose, 11);
    const bool has_right_hip = IsValidKeypoint(pose, 12);
    if (has_left_shoulder && has_right_shoulder && has_left_hip && has_right_hip) {
        const PosePoint2f& ls = pose.keypoints[5].point;
        const PosePoint2f& rs = pose.keypoints[6].point;
        const PosePoint2f& lh = pose.keypoints[11].point;
        const PosePoint2f& rh = pose.keypoints[12].point;
        const float shoulder_center_y = (ls.y + rs.y) * 0.5f;
        const float hip_center_y = (lh.y + rh.y) * 0.5f;
        sample.torso_center_y = (shoulder_center_y + hip_center_y) * 0.5f;
        const float min_x = std::min(std::min(ls.x, rs.x), std::min(lh.x, rh.x));
        const float max_x = std::max(std::max(ls.x, rs.x), std::max(lh.x, rh.x));
        const float min_y = std::min(std::min(ls.y, rs.y), std::min(lh.y, rh.y));
        const float max_y = std::max(std::max(ls.y, rs.y), std::max(lh.y, rh.y));
        const float horizontal_span = std::fabs(max_x - min_x);
        const float vertical_span = std::fabs(max_y - min_y);
        sample.upright_ratio = vertical_span > 1e-3f ? (horizontal_span / vertical_span) : std::numeric_limits<float>::infinity();
    }

    if (IsValidKeypoint(pose, 9)) {
        sample.left_wrist_x = pose.keypoints[9].point.x;
        sample.left_wrist_y = pose.keypoints[9].point.y;
        sample.has_left_wrist = true;
    }
    if (IsValidKeypoint(pose, 10)) {
        sample.right_wrist_x = pose.keypoints[10].point.x;
        sample.right_wrist_y = pose.keypoints[10].point.y;
        sample.has_right_wrist = true;
    }
    return sample;
}

static PoseTrackSample MatchPoseToDetection(const Rect& bbox, const std::shared_ptr<PoseResult>& pose_result) {
    if (!pose_result) return {};

    float best_iou = 0.0f;
    const PoseItem* best_pose = nullptr;
    for (const auto& pose : pose_result->items) {
        const float iou = BboxIoU(bbox, pose.bbox);
        if (iou > best_iou) {
            best_iou = iou;
            best_pose = &pose;
        }
    }
    if (!best_pose || best_iou <= 0.0f) return {};
    return BuildPoseTrackSample(*best_pose);
}

static PoseTrackSample MatchPoseToTrackId(int track_id, const Rect& bbox, const std::shared_ptr<PoseResult>& pose_result) {
    if (!pose_result) return {};

    const PoseItem* best_pose = nullptr;
    float best_iou = 0.0f;
    for (const auto& pose : pose_result->items) {
        if (pose.track_id != track_id) continue;
        const float iou = BboxIoU(bbox, pose.bbox);
        if (iou >= best_iou) {
            best_iou = iou;
            best_pose = &pose;
        }
    }
    if (best_pose) return BuildPoseTrackSample(*best_pose);
    return MatchPoseToDetection(bbox, pose_result);
}

static float WristMotion(const PoseTrackSample& newer, const PoseTrackSample& older) {
    float motion = 0.0f;
    if (newer.has_left_wrist && older.has_left_wrist) {
        const float dx = newer.left_wrist_x - older.left_wrist_x;
        const float dy = newer.left_wrist_y - older.left_wrist_y;
        motion += std::sqrt(dx * dx + dy * dy);
    }
    if (newer.has_right_wrist && older.has_right_wrist) {
        const float dx = newer.right_wrist_x - older.right_wrist_x;
        const float dy = newer.right_wrist_y - older.right_wrist_y;
        motion += std::sqrt(dx * dx + dy * dy);
    }
    return motion;
}

static float MinimumWristDistance(const PoseTrackSample& lhs, const PoseTrackSample& rhs) {
    float best = std::numeric_limits<float>::infinity();
    auto update = [&](bool has_a, float ax, float ay, bool has_b, float bx, float by) {
        if (!has_a || !has_b) return;
        const float dx = ax - bx;
        const float dy = ay - by;
        best = std::min(best, std::sqrt(dx * dx + dy * dy));
    };
    update(lhs.has_left_wrist, lhs.left_wrist_x, lhs.left_wrist_y,
           rhs.has_left_wrist, rhs.left_wrist_x, rhs.left_wrist_y);
    update(lhs.has_left_wrist, lhs.left_wrist_x, lhs.left_wrist_y,
           rhs.has_right_wrist, rhs.right_wrist_x, rhs.right_wrist_y);
    update(lhs.has_right_wrist, lhs.right_wrist_x, lhs.right_wrist_y,
           rhs.has_left_wrist, rhs.left_wrist_x, rhs.left_wrist_y);
    update(lhs.has_right_wrist, lhs.right_wrist_x, lhs.right_wrist_y,
           rhs.has_right_wrist, rhs.right_wrist_x, rhs.right_wrist_y);
    return std::isfinite(best) ? best : std::numeric_limits<float>::infinity();
}

static BehaviorEventType ToBehaviorEventType(ActionEventKind kind) {
    return kind == ActionEventKind::Fall ? BehaviorEventType::Fall : BehaviorEventType::Fight;
}

}  // namespace

struct ActionRecogNode::Impl {
    std::string init_err;
    std::vector<ActionEventRule> rules;
    std::map<int, std::deque<TrackSample>> history;
};

ActionRecogNode::ActionRecogNode() : impl_(std::make_unique<Impl>()) {}

ActionRecogNode::~ActionRecogNode() = default;

std::string ActionRecogNode::Id() const {
    return id_;
}

std::string ActionRecogNode::Type() const {
    return "action_recog";
}

bool ActionRecogNode::Init(const SimpleJson& config, const NodeContext& ctx) {
    id_ = config.ValueOr<std::string>("id", "action_recog");
    if (!ParseRules(config, impl_->rules, impl_->init_err)) {
        LogError("[action_recog] invalid config: " + impl_->init_err);
        return false;
    }
    output_queues_ = ctx.output_queues;
    return true;
}

bool ActionRecogNode::Start() {
    return true;
}

void ActionRecogNode::Stop() {}

NodeStatus ActionRecogNode::Process(FramePtr frame) {
    if (!frame) return NodeStatus::DROP;

    EnsureBehaviorEvents(*frame);
    if (!frame->det) {
        PushToDownstream(frame);
        return NodeStatus::OK;
    }

    std::map<int, Rect> current_tracks;
    for (const auto& det : frame->det->items) {
        if (det.track_id < 0) continue;
        current_tracks[det.track_id] = det.bbox;
        auto& history = impl_->history[det.track_id];
        history.push_back(TrackSample{frame->pts, det.bbox, MatchPoseToTrackId(det.track_id, det.bbox, frame->pose)});
        while (!history.empty() && frame->pts > history.front().pts &&
               (frame->pts - history.front().pts) > impl_->rules.front().window_ms) {
            history.pop_front();
        }
    }

    for (const auto& rule : impl_->rules) {
        if (rule.kind == ActionEventKind::Fall) {
            for (const auto& [track_id, bbox] : current_tracks) {
                auto it = impl_->history.find(track_id);
                if (it == impl_->history.end() || it->second.empty()) continue;
                const auto& first = it->second.front();
                const float drop = CenterY(bbox) - CenterY(first.bbox);
                const float aspect_delta = AspectRatio(bbox) - AspectRatio(first.bbox);
                const uint64_t duration = frame->pts >= first.pts ? (frame->pts - first.pts) : 0;
                const bool bbox_enabled = rule.fall_bbox.enabled;
                const bool pose_enabled = rule.fall_pose.enabled;
                const bool bbox_trigger = drop >= rule.fall_bbox.min_drop_pixels &&
                                          aspect_delta >= rule.fall_bbox.min_aspect_ratio_delta;
                bool pose_trigger = false;
                if (pose_enabled &&
                    first.pose.matched && it->second.back().pose.matched) {
                    const float torso_drop = it->second.back().pose.torso_center_y - first.pose.torso_center_y;
                    pose_trigger = torso_drop >= rule.fall_pose.min_torso_drop_pixels &&
                                   it->second.back().pose.upright_ratio >= rule.fall_pose.max_upright_ratio;
                }
                if (!EvaluateSignalMode(rule.signal_mode, bbox_trigger, pose_trigger, bbox_enabled, pose_enabled)) continue;
                if (duration < rule.activate_duration_ms) continue;

                BehaviorEventItem event;
                event.type = ToBehaviorEventType(rule.kind);
                event.status = BehaviorEventStatus::Active;
                event.score = 1.0f;
                event.bbox = bbox;
                event.track_ids.push_back(track_id);
                event.start_pts = first.pts;
                event.last_pts = frame->pts;
                event.duration_ms = duration;
                event.source = "action_recog";
                frame->behavior_events->items.push_back(std::move(event));
            }
        } else {
            std::set<std::pair<int, int>> emitted_pairs;
            for (auto left = current_tracks.begin(); left != current_tracks.end(); ++left) {
                for (auto right = std::next(left); right != current_tracks.end(); ++right) {
                    const float proximity = CenterDistance(left->second, right->second);
                    if (proximity > rule.fight_bbox.proximity_pixels) continue;

                    const auto it_left = impl_->history.find(left->first);
                    const auto it_right = impl_->history.find(right->first);
                    if (it_left == impl_->history.end() || it_right == impl_->history.end()) continue;
                    if (it_left->second.empty() || it_right->second.empty()) continue;

                    const float left_motion = MotionDistance(left->second, it_left->second.front().bbox);
                    const float right_motion = MotionDistance(right->second, it_right->second.front().bbox);
                    const uint64_t duration_left = frame->pts >= it_left->second.front().pts ? (frame->pts - it_left->second.front().pts) : 0;
                    const uint64_t duration_right = frame->pts >= it_right->second.front().pts ? (frame->pts - it_right->second.front().pts) : 0;
                    const uint64_t duration = std::min(duration_left, duration_right);

                    const float combined_motion = left_motion + right_motion;
                    const bool bbox_enabled = rule.fight_bbox.enabled;
                    const bool pose_enabled = rule.fight_pose.enabled;
                    const bool bbox_trigger = combined_motion >= rule.fight_bbox.min_motion_pixels;
                    bool pose_trigger = false;
                    const auto& left_pose_now = it_left->second.back().pose;
                    const auto& right_pose_now = it_right->second.back().pose;
                    if (pose_enabled &&
                        it_left->second.front().pose.matched && it_right->second.front().pose.matched &&
                        left_pose_now.matched && right_pose_now.matched) {
                        const float pose_motion = WristMotion(left_pose_now, it_left->second.front().pose) +
                                                  WristMotion(right_pose_now, it_right->second.front().pose);
                        const float wrist_distance = MinimumWristDistance(left_pose_now, right_pose_now);
                        pose_trigger = pose_motion >= rule.fight_pose.min_wrist_motion_pixels &&
                                       wrist_distance <= rule.fight_pose.max_wrist_distance_pixels;
                    }
                    if (!EvaluateSignalMode(rule.signal_mode, bbox_trigger, pose_trigger, bbox_enabled, pose_enabled)) continue;
                    if (duration < rule.activate_duration_ms) continue;

                    const auto pair_key = std::make_pair(left->first, right->first);
                    if (!emitted_pairs.insert(pair_key).second) continue;

                    BehaviorEventItem event;
                    event.type = ToBehaviorEventType(rule.kind);
                    event.status = BehaviorEventStatus::Active;
                    event.score = 1.0f;
                    event.bbox = Rect{
                        std::min(left->second.x, right->second.x),
                        std::min(left->second.y, right->second.y),
                        std::max(left->second.x + left->second.w, right->second.x + right->second.w) - std::min(left->second.x, right->second.x),
                        std::max(left->second.y + left->second.h, right->second.y + right->second.h) - std::min(left->second.y, right->second.y)
                    };
                    event.track_ids.push_back(left->first);
                    event.track_ids.push_back(right->first);
                    event.start_pts = frame->pts - duration;
                    event.last_pts = frame->pts;
                    event.duration_ms = duration;
                    event.source = "action_recog";
                    frame->behavior_events->items.push_back(std::move(event));
                }
            }
        }
    }

    PushToDownstream(frame);
    return NodeStatus::OK;
}

void ActionRecogNode::EnsureBehaviorEvents(Frame& frame) {
    if (!frame.behavior_events) {
        frame.behavior_events = std::make_shared<BehaviorEventResult>();
    }
}

void ActionRecogNode::PushToDownstream(const FramePtr& frame) {
    for (auto& q : output_queues_) {
        if (q) q->Push(frame);
    }
}

#ifndef RK3588_TEST_BUILD
REGISTER_NODE(ActionRecogNode, "action_recog");
#endif

}  // namespace rk3588