#pragma once

#include "pipeline_types.hpp"
#include "trt_engine.hpp"

#include <opencv2/core.hpp>

#include <deque>
#include <array>
#include <optional>
#include <string>
#include <unordered_map>

namespace fangjia {

class BoxTracker {
public:
    std::vector<int> update(std::vector<Box>& boxes, int frameIndex, int width, int height);

private:
    struct Track { Box box; int lastFrame{}; };
    int nextId_{1};
    std::unordered_map<int, Track> tracks_;
};

class RtmwEngine {
public:
    RtmwEngine(
        std::filesystem::path const& enginePath,
        TrtLogger& logger,
        int maxSourceWidth,
        int maxSourceHeight,
        int maxPeople = 3);
    ~RtmwEngine();

    RtmwEngine(RtmwEngine const&) = delete;
    RtmwEngine& operator=(RtmwEngine const&) = delete;

    std::vector<PersonPose> infer(
        cv::Mat const& frame,
        std::vector<Box> const& boxes,
        cudaStream_t stream,
        float* preprocessMs = nullptr,
        float* inferenceMs = nullptr,
        float* postprocessMs = nullptr);

    void warmup(cudaStream_t stream);

private:
    TrtEngine engine_;
    int maxSourceWidth_{};
    int maxSourceHeight_{};
    int maxPeople_{};
    std::uint8_t* sourceDevice_{};
    std::uint8_t* cropsDevice_{};
    float* networkDevice_{};
    std::vector<double> centersHost_;
    std::vector<double> scalesHost_;
    std::vector<float> simccXHost_;
    std::vector<float> simccYHost_;
};

class PoseFusionTracker {
public:
    PoseFrameResult update(
        std::vector<PersonPose> const& raw,
        std::vector<PersonPose> const& isolated,
        int frameIndex);

private:
    struct Vec2 { double x{}; double y{}; };
    struct SideFusion {
        std::array<Vec2, 3> limbPrevious{};
        std::array<Vec2, 3> limbVelocity{};
        std::array<Vec2, 21> handPrevious{};
        std::array<Vec2, 21> handVelocity{};
        bool hasLimbPrevious{};
        bool hasLimbVelocity{};
        bool hasHandPrevious{};
        bool hasHandVelocity{};
        std::string limbSource;
        std::string handSource;
    };
    struct Endpoint {
        bool hasWrist{};
        Vec2 wrist{};
        Vec2 velocity{};
        std::array<Vec2, 21> handShape{};
        bool hasHandShape{};
        int lastObservedFrame{-1};
        int missedFrames{};
        bool hasLastElbow{};
        Vec2 lastElbow{};
        bool hasLastDirection{};
        Vec2 lastDirection{};
    };
    struct PersonState {
        std::array<SideFusion, 2> fusion;
        std::array<Endpoint, 2> endpoints;
        std::array<std::deque<std::array<double, 2>>, 2> boneSamples;
        std::array<Vec2, 133> smoothed{};
        bool hasSmoothed{};
        int lastFrame{-1};
    };
    std::unordered_map<int, PersonState> states_;
};

} // namespace fangjia