#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
#include <pybind11/numpy.h>

#include "particri/math.hpp"
#include "particri/particle.hpp"
#define FNL_IMPL
#include "particri/forces.hpp"
#include "particri/emitter.hpp"
#include "particri/system.hpp"

namespace py = pybind11;
using namespace particri;

PYBIND11_MODULE(_particri_core, m) {
    m.doc() = "Particri C++ High-Performance Particle Engine Core";

    // Enums
    py::enum_<CurveType>(m, "CurveType")
        .value("Linear", CurveType::Linear)
        .value("Bell", CurveType::Bell)
        .value("EaseIn", CurveType::EaseIn)
        .value("EaseOut", CurveType::EaseOut)
        .value("GrowAndFade", CurveType::GrowAndFade)
        .value("Pulse", CurveType::Pulse)
        .value("Decay", CurveType::Decay)
        .export_values();

    py::enum_<ParticleShape>(m, "ParticleShape")
        .value("Bokeh", ParticleShape::Bokeh)
        .value("Star4", ParticleShape::Star4)
        .value("Star8", ParticleShape::Star8)
        .value("Halo", ParticleShape::Halo)
        .value("Streak", ParticleShape::Streak)
        .value("Glitter", ParticleShape::Glitter)
        .value("Custom", ParticleShape::Custom)
        .export_values();

    py::enum_<BlendMode>(m, "BlendMode")
        .value("Additive", BlendMode::Additive)
        .value("Screen", BlendMode::Screen)
        .value("Alpha", BlendMode::Alpha)
        .value("ColorDodge", BlendMode::ColorDodge)
        .export_values();

    py::enum_<EmitterShape>(m, "EmitterShape")
        .value("Point", EmitterShape::Point)
        .value("Circle", EmitterShape::Circle)
        .value("Ring", EmitterShape::Ring)
        .value("Box", EmitterShape::Box)
        .value("Line", EmitterShape::Line)
        .value("LayerMask", EmitterShape::LayerMask)
        .export_values();

    // Math structs
    py::class_<Vec2>(m, "Vec2")
        .def(py::init<float, float>(), py::arg("x") = 0.0f, py::arg("y") = 0.0f)
        .def_readwrite("x", &Vec2::x)
        .def_readwrite("y", &Vec2::y)
        .def("__repr__", [](const Vec2& v) {
            return "Vec2(" + std::to_string(v.x) + ", " + std::to_string(v.y) + ")";
        });

    py::class_<Color4>(m, "Color4")
        .def(py::init<float, float, float, float>(),
             py::arg("r") = 1.0f, py::arg("g") = 1.0f, py::arg("b") = 1.0f, py::arg("a") = 1.0f)
        .def_readwrite("r", &Color4::r)
        .def_readwrite("g", &Color4::g)
        .def_readwrite("b", &Color4::b)
        .def_readwrite("a", &Color4::a)
        .def_static("from_hex", &Color4::from_hex, py::arg("hex"), py::arg("alpha") = 1.0f)
        .def("__repr__", [](const Color4& c) {
            return "Color4(" + std::to_string(c.r) + ", " + std::to_string(c.g) + ", " +
                   std::to_string(c.b) + ", " + std::to_string(c.a) + ")";
        });

    // EmitterConfig
    py::class_<EmitterConfig>(m, "EmitterConfig")
        .def(py::init<>())
        .def_readwrite("shape", &EmitterConfig::shape)
        .def_readwrite("position", &EmitterConfig::position)
        .def_readwrite("size", &EmitterConfig::size)
        .def_readwrite("line_end", &EmitterConfig::line_end)
        .def_readwrite("rate", &EmitterConfig::rate)
        .def_readwrite("speed_min", &EmitterConfig::speed_min)
        .def_readwrite("speed_max", &EmitterConfig::speed_max)
        .def_readwrite("angle_min", &EmitterConfig::angle_min)
        .def_readwrite("angle_max", &EmitterConfig::angle_max)
        .def_readwrite("lifetime_min", &EmitterConfig::lifetime_min)
        .def_readwrite("lifetime_max", &EmitterConfig::lifetime_max)
        .def_readwrite("size_start_min", &EmitterConfig::size_start_min)
        .def_readwrite("size_start_max", &EmitterConfig::size_start_max)
        .def_readwrite("size_end_min", &EmitterConfig::size_end_min)
        .def_readwrite("size_end_max", &EmitterConfig::size_end_max)
        .def_readwrite("color_start", &EmitterConfig::color_start)
        .def_readwrite("color_start_variance", &EmitterConfig::color_start_variance)
        .def_readwrite("color_end", &EmitterConfig::color_end)
        .def_readwrite("particle_shape", &EmitterConfig::particle_shape)
        .def_readwrite("size_curve", &EmitterConfig::size_curve)
        .def_readwrite("opacity_curve", &EmitterConfig::opacity_curve)
        .def_readwrite("inherit_layer_color", &EmitterConfig::inherit_layer_color)
        .def_readwrite("rotation_speed_min", &EmitterConfig::rotation_speed_min)
        .def_readwrite("rotation_speed_max", &EmitterConfig::rotation_speed_max)
        .def_readwrite("twinkle_speed_min", &EmitterConfig::twinkle_speed_min)
        .def_readwrite("twinkle_speed_max", &EmitterConfig::twinkle_speed_max)
        .def_readwrite("drag", &EmitterConfig::drag)
        .def_readwrite("radial_velocity", &EmitterConfig::radial_velocity);

    // AuxConfig (Trapcode Particular Auxiliary System)
    py::class_<AuxConfig>(m, "AuxConfig")
        .def(py::init<>())
        .def_readwrite("enabled", &AuxConfig::enabled)
        .def_readwrite("rate", &AuxConfig::rate)
        .def_readwrite("lifetime", &AuxConfig::lifetime)
        .def_readwrite("size_start", &AuxConfig::size_start)
        .def_readwrite("size_end", &AuxConfig::size_end)
        .def_readwrite("color_start", &AuxConfig::color_start)
        .def_readwrite("color_end", &AuxConfig::color_end)
        .def_readwrite("inherit_color", &AuxConfig::inherit_color)
        .def_readwrite("velocity_inheritance", &AuxConfig::velocity_inheritance)
        .def_readwrite("spread", &AuxConfig::spread)
        .def_readwrite("shape", &AuxConfig::shape)
        .def_readwrite("size_curve", &AuxConfig::size_curve)
        .def_readwrite("opacity_curve", &AuxConfig::opacity_curve)
        .def_readwrite("emit_at_death", &AuxConfig::emit_at_death);

    // ParticleSystem
    py::class_<ParticleSystem>(m, "ParticleSystem")
        .def(py::init<size_t, int, int>(),
             py::arg("max_particles") = 20000,
             py::arg("width") = 1920,
             py::arg("height") = 1080)
        .def("add_emitter", &ParticleSystem::add_emitter, py::arg("config"), py::arg("seed") = 42)
        .def("add_gravity", [](ParticleSystem& s, float gx, float gy) {
            s.add_force(std::make_shared<GravityForce>(Vec2(gx, gy)));
        }, py::arg("gx"), py::arg("gy"))
        .def("add_drag", [](ParticleSystem& s, float coeff) {
            s.add_force(std::make_shared<DragForce>(coeff));
        }, py::arg("coeff"))
        .def("add_vortex", [](ParticleSystem& s, float cx, float cy, float strength, float radial_pull, float radius) {
            s.add_force(std::make_shared<VortexForce>(Vec2(cx, cy), strength, radial_pull, radius));
        }, py::arg("cx"), py::arg("cy"), py::arg("strength"), py::arg("radial_pull") = 0.0f, py::arg("radius") = 800.0f)
        .def("add_attractor", [](ParticleSystem& s, float tx, float ty, float strength, float min_d, float max_d) {
            s.add_force(std::make_shared<PointAttractorForce>(Vec2(tx, ty), strength, min_d, max_d));
        }, py::arg("tx"), py::arg("ty"), py::arg("strength"), py::arg("min_distance") = 20.0f, py::arg("max_distance") = 800.0f)
        .def("add_curl_noise", [](ParticleSystem& s, int seed, float strength, float frequency, float speed) {
            s.add_force(std::make_shared<CurlNoiseForce>(seed, strength, frequency, speed));
        }, py::arg("seed") = 1337, py::arg("strength") = 100.0f, py::arg("frequency") = 0.003f, py::arg("speed") = 0.3f)
        .def("add_brownian", [](ParticleSystem& s, float strength, uint32_t seed) {
            s.add_force(std::make_shared<BrownianForce>(strength, seed));
        }, py::arg("strength") = 30.0f, py::arg("seed") = 42)
        .def("add_bounce_plane", [](ParticleSystem& s, float coord, float bounce, float friction, bool is_horizontal, bool is_max) {
            s.add_force(std::make_shared<BouncePlaneForce>(coord, bounce, friction, is_horizontal, is_max));
        }, py::arg("coord") = 900.0f, py::arg("bounce") = 0.65f, py::arg("friction") = 0.85f, py::arg("is_horizontal") = true, py::arg("is_max") = true)
        .def("set_aux_config", &ParticleSystem::set_aux_config, py::arg("config"))
        .def("get_aux_config", [](ParticleSystem& s) { return s.aux_config(); })
        .def("set_emitter_layer_mask", [](ParticleSystem& s, size_t emitter_idx, py::array_t<float> img_arr, float threshold) {
            auto em = s.get_emitter(emitter_idx);
            if (!em) return;
            auto buf = img_arr.request();
            if (buf.ndim != 3 || buf.shape[2] != 4) {
                throw std::runtime_error("Layer mask array must have shape (H, W, 4) float32");
            }
            int h = static_cast<int>(buf.shape[0]);
            int w = static_cast<int>(buf.shape[1]);
            const float* ptr = static_cast<const float*>(buf.ptr);
            em->set_layer_mask(ptr, w, h, threshold);
        }, py::arg("emitter_idx"), py::arg("image_data"), py::arg("threshold") = 0.05f)
        .def("clear_forces", &ParticleSystem::clear_forces)
        .def("clear_emitters", &ParticleSystem::clear_emitters)
        .def("burst", &ParticleSystem::burst, py::arg("emitter_index"), py::arg("count"))
        .def("update", &ParticleSystem::update, py::arg("dt"))
        .def("render", &ParticleSystem::render,
             py::arg("blend_mode") = BlendMode::Additive,
             py::arg("enable_bloom") = true,
             py::arg("bloom_threshold") = 0.35f,
             py::arg("bloom_intensity") = 0.75f,
             py::arg("bloom_radius") = 8)
        .def("resize", &ParticleSystem::resize, py::arg("width"), py::arg("height"))
        .def_property_readonly("width", &ParticleSystem::width)
        .def_property_readonly("height", &ParticleSystem::height)
        .def_property_readonly("active_count", &ParticleSystem::active_count)
        .def("clear_particles", &ParticleSystem::clear_particles)
        .def("get_image", [](ParticleSystem& s) {
            int w = s.width();
            int h = s.height();
            py::array_t<uint8_t> result({h, w, 4});
            auto buf = result.request();
            uint8_t* ptr = static_cast<uint8_t*>(buf.ptr);
            s.framebuffer().to_uint8_rgba(ptr);
            return result;
        }, "Returns rendered frame as an (H, W, 4) uint8 RGBA NumPy array")
        .def("get_hdr_image", [](ParticleSystem& s) {
            int w = s.width();
            int h = s.height();
            py::array_t<float> result({h, w, 4});
            auto buf = result.request();
            float* ptr = static_cast<float*>(buf.ptr);
            std::memcpy(ptr, s.framebuffer().data.data(), w * h * 4 * sizeof(float));
            return result;
        }, "Returns rendered frame as an (H, W, 4) float32 HDR RGBA NumPy array");
}