#pragma once

#include <vector>
#include <variant>

#include <tbb/tbb.h>
#include <utils/eigen_alias.hpp>
#include <macros.h>

#include "primitive_descriptor.h"

// =========================================================================================================================

class evaluation_routine_tag
{
};

class closest_point_routine_tag
{
};

template <typename T>
static constexpr bool is_process_routine_tag_v =
    std::is_same_v<T, evaluation_routine_tag> || std::is_same_v<T, closest_point_routine_tag>;

template <typename T>
static constexpr bool is_evaluation_routine_v = std::is_same_v<T, evaluation_routine_tag>;

template <typename T>
static constexpr bool is_closest_point_routine_v = std::is_same_v<T, closest_point_routine_tag>;

template <typename T>
using process_return_type_t = std::conditional_t<is_evaluation_routine_v<T>, double, Eigen::Vector3d>;

// static constexpr evaluation_routine_tag    evaluation_tag{};
// static constexpr closest_point_routine_tag closest_point_tag{};

// =========================================================================================================================

template <size_t Dim = 3>
using aabb_t = std::conditional_t<Dim == 2, Eigen::AlignedBox2d, Eigen::AlignedBox3d>;

using legal_primitive_descriptor_t = std::variant<constant_descriptor_t,
                                                  plane_descriptor_t,
                                                  sphere_descriptor_t,
                                                  cylinder_descriptor_t,
                                                  cone_descriptor_t,
                                                  box_descriptor_t,
                                                  mesh_descriptor_t,
                                                  extrude_polyline_descriptor_t,
                                                  extrude_helixline_descriptor_t>;

struct PE_API primitive_node_t {
    aabb_t<>       aabb{};
    void*          desc{nullptr}; // Type conversion when using
    primitive_type type{};

    primitive_node_t(const legal_primitive_descriptor_t&) noexcept;
    primitive_node_t(legal_primitive_descriptor_t&&) noexcept;
    ~primitive_node_t() noexcept;

    double          evaluate_sdf([[maybe_unused]] const Eigen::Ref<const Eigen::Vector3d>&) const;
    Eigen::Vector3d evaluate_cpm([[maybe_unused]] const Eigen::Ref<const Eigen::Vector3d>&) const;
};

static constexpr auto EPSILON     = std::numeric_limits<double>::epsilon() * 1e6;
static constexpr auto PI          = 3.14159265358979323846;
static constexpr auto TWO_PI      = PI * 2;
static constexpr auto PI2         = PI / 2;
static constexpr auto INV_PI      = 1. / PI;
static constexpr auto INV_TWO_PI  = 1. / TWO_PI;
static const auto     x_direction = Eigen::Vector3d{1.0, 0.0, 0.0};

inline void vec3d_conversion(const raw_vector3d_t& src, Eigen::Ref<Eigen::Vector3d> dst)
{
    dst = Eigen::Map<const Eigen::Vector3d>(&src.x);
}

inline void vec3d_conversion(raw_vector3d_t&& src, Eigen::Ref<Eigen::Vector3d> dst)
{
    std::move(&src.x, &src.x + 3, dst.data());
}

inline double sign(const double t) { return t >= 0.0 ? 1.0 : -1.0; }

// Eigen has a type Isometry which supports similar operations as this
// but we use AffineCompact to get lower storage cost
// so we have to implement a helper function to apply the operation of Isometry to AffineCompact
inline auto inversed_affine_transform(const Eigen::Transform<double, 3, Eigen::AffineCompact>& trs)
{
    Eigen::Transform<double, 3, Eigen::AffineCompact> result;

    auto linear_part                                = result.matrix().template topLeftCorner<3, 3>();
    linear_part                                     = trs.linear().transpose();
    result.matrix().template topRightCorner<3, 1>() = -linear_part * trs.translation();
    result.makeAffine();

    return result;
}

namespace internal
{
template <typename T>
using stl_vector = std::vector<T, tbb::tbb_allocator<T>>;

struct line_closest_param_t {
    Eigen::Vector3d point{};
    double          t{};
    bool            is_peak_value{};
};

#define DEF_PRIMITIVE_COMMON_METHODS(internal_type, descriptor_type)                               \
    internal_type(const descriptor_type&, aabb_t<>&);                                              \
    internal_type(descriptor_type&&, aabb_t<>&);                                                   \
    double          evaluate_sdf([[maybe_unused]] const Eigen::Ref<const Eigen::Vector3d>&) const; \
    Eigen::Vector3d evaluate_cpm([[maybe_unused]] const Eigen::Ref<const Eigen::Vector3d>&) const;

struct PE_API constant {
    double value{};

    DEF_PRIMITIVE_COMMON_METHODS(constant, constant_descriptor_t)
};

struct PE_API plane {
    Eigen::Vector3d point{};
    Eigen::Vector3d normal{};

    DEF_PRIMITIVE_COMMON_METHODS(plane, plane_descriptor_t)
};

struct PE_API sphere {
    Eigen::Vector3d center{};
    double          radius{};

    DEF_PRIMITIVE_COMMON_METHODS(sphere, sphere_descriptor_t)
};

struct PE_API cylinder {
    Eigen::Vector3d bottom_center{};
    Eigen::Vector3d offset{};
    double          radius{};

    DEF_PRIMITIVE_COMMON_METHODS(cylinder, cylinder_descriptor_t)
};

struct PE_API cone {
    Eigen::Vector3d top_center{};
    Eigen::Vector3d bottom_center{};
    double          radius1{};
    double          radius2{};

    DEF_PRIMITIVE_COMMON_METHODS(cone, cone_descriptor_t)
};

struct PE_API box {
    Eigen::Vector3d center{};
    Eigen::Vector3d half_size{};

    DEF_PRIMITIVE_COMMON_METHODS(box, box_descriptor_t)
};

struct PE_API mesh {
    stl_vector<Eigen::Vector3d>      vertices{};
    stl_vector<stl_vector<uint32_t>> faces{};

    DEF_PRIMITIVE_COMMON_METHODS(mesh, mesh_descriptor_t)
};

// CAUTION: in polyline local space, the X/Y axis should be according to local N/B directions
struct polyline {
    std::vector<Eigen::Vector2d> vertices{}; //
                                             // for a straight line, it consists of two vertices; for a circle arc, it consists
                                             // of three vertices: start point, circle center, and a point on the circle to
                                             // construct local XY coordinate system
    std::vector<double>          thetas{};   // for a straight line, this should be 0
    std::vector<uint8_t>         start_indices{}; // has the same size as thetas, indicating the starting index of each segment

    void build_as_axis(const polyline_descriptor_t&,
                       const raw_vector3d_t&,
                       Eigen::Transform<double, 3, Eigen::AffineCompact>&,
                       aabb_t<>&);
    void build_as_axis(polyline_descriptor_t&&,
                       const raw_vector3d_t&,
                       Eigen::Transform<double, 3, Eigen::AffineCompact>&,
                       aabb_t<>&);
    void build_as_profile(const polyline_descriptor_t&,
                          const Eigen::Ref<const Eigen::Vector3d>&,
                          const Eigen::Ref<const Eigen::Vector3d>&,
                          const Eigen::Ref<const Eigen::Vector3d>&,
                          aabb_t<2>&);
    void build_as_profile(polyline_descriptor_t&&,
                          const Eigen::Ref<const Eigen::Vector3d>&,
                          const Eigen::Ref<const Eigen::Vector3d>&,
                          const Eigen::Ref<const Eigen::Vector3d>&,
                          aabb_t<2>&);

    [[nodiscard]] Eigen::Vector2d                         calculate_tangent(double) const;
    [[nodiscard]] Eigen::Vector2d                         calculate_normal(double) const;
    // CAUTION: make sure the input point is in the local space of the polyline
    // return type: [line_closest_param_t, SDF]
    [[nodiscard]] std::pair<line_closest_param_t, double> calculate_closest_param(
        const Eigen::Ref<const Eigen::Vector3d>&) const;
};

struct helixline {
    double radius{};
    double total_theta{};
    double height{};

    // CAUTION: here returned aabb is in helixline's local space
    helixline(const helixline_descriptor_t&, Eigen::Transform<double, 3, Eigen::AffineCompact>&, aabb_t<>&);
    helixline(helixline_descriptor_t&&, Eigen::Transform<double, 3, Eigen::AffineCompact>&, aabb_t<>&);

    [[nodiscard]] Eigen::Vector3d      calculate_tangent(double) const;
    [[nodiscard]] Eigen::Vector3d      calculate_normal(double) const;
    // CAUTION: make sure the input point is in the local space of the helixline
    [[nodiscard]] line_closest_param_t calculate_closest_param(const Eigen::Ref<const Eigen::Vector3d>&) const;
};

struct PE_API extrude_polyline {
    Eigen::Transform<double, 3, Eigen::AffineCompact> axis_to_world{};
    polyline                                          axis;
    polyline                                          profile;

    DEF_PRIMITIVE_COMMON_METHODS(extrude_polyline, extrude_polyline_descriptor_t)
};

struct PE_API extrude_helixline {
    Eigen::Transform<double, 3, Eigen::AffineCompact> axis_to_world{};
    helixline                                         axis;
    polyline                                          profile;

    DEF_PRIMITIVE_COMMON_METHODS(extrude_helixline, extrude_helixline_descriptor_t)
};

#undef DEF_PRIMITIVE_COMMON_METHODS
} // namespace internal