#include <torch/extension.h>

#include "input_metadata/prepare_tensor_input.hpp"
#include "outputs.hpp"
#include "prepare_sample.hpp"

namespace deepinfer::engine {

void pybind_engine_output(pybind11::module &engine) {
// def_readonly plus a matching static `get_<field>(ptr)` accessor used by the
// raw-pointer fast path on the Python side.
#define DEF_OUTPUT_FIELD(name)                                                                                                             \
    def_readonly(#name, &RequestOutput::name)                                                                                              \
        .def_static("get_" #name, [](uint64_t ptr) { return reinterpret_cast<RequestOutput *>(ptr)->name; }, py::arg("ptr"))

    py::class_<RequestOutput>(engine, "RequestOutput")
        .def_static(
            "clone", [](uint64_t ptr) { return std::make_unique<RequestOutput>(*reinterpret_cast<RequestOutput *>(ptr)); }, py::arg("ptr"))
        .DEF_OUTPUT_FIELD(request_id)
        .DEF_OUTPUT_FIELD(prompt_len)
        .DEF_OUTPUT_FIELD(token_id)
        .DEF_OUTPUT_FIELD(finish_reason)
        .DEF_OUTPUT_FIELD(finished)
        .def_readonly("max_num_batched_tokens", &RequestOutput::max_num_batched_tokens)
        .DEF_OUTPUT_FIELD(token_logprobs)
        .DEF_OUTPUT_FIELD(nextn_logprobs)
        .DEF_OUTPUT_FIELD(sum_smaller_probs)
        .DEF_OUTPUT_FIELD(top_logprobs)
        .DEF_OUTPUT_FIELD(top_token_ids)
        .DEF_OUTPUT_FIELD(prompt_logprobs)
        .DEF_OUTPUT_FIELD(num_output_tokens)
        .DEF_OUTPUT_FIELD(draft_tok)
        .DEF_OUTPUT_FIELD(accept_tok)
        .DEF_OUTPUT_FIELD(confidence_score);
#undef DEF_OUTPUT_FIELD

    py::class_<RequestOutputPool>(engine, "RequestOutputPool")
        .def(py::init<>())
        .def("release", &RequestOutputPool::release, py::arg("unused_request_outputs"));

    engine.def("build_request_outputs", &build_request_outputs, py::arg("scheduler_outputs"), py::arg("max_num_batched_tokens"),
        py::arg("token_ids"), py::arg("token_logprobs"), py::arg("top_logprobs_"), py::arg("top_token_ids_"), py::arg("sum_smaller_probs_"),
        py::arg("accepted_token_ids_"), py::arg("accepted_token_logprobs_"), py::arg("accepted_top_logprobs_"),
        py::arg("accepted_top_token_ids_"), py::arg("accepted_sum_smaller_probs_"), py::arg("accepted_token_last_logprobs_"),
        py::arg("pushed_tokens_before_finish"), py::arg("prompt_logprobs"), py::arg("request_output_pool"),
        py::arg("first_token_positions_") = py::none());

    py::enum_<YocoPrefillTruncationMode>(engine, "YocoPrefillTruncationMode")
        .value("OFF", YocoPrefillTruncationMode::OFF)
        .value("MASK", YocoPrefillTruncationMode::MASK)
        .value("COMPACT", YocoPrefillTruncationMode::COMPACT);

    py::class_<PrepareArgs>(engine, "PrepareArgs")
        .def(py::init<>())
        .def_readwrite("seq_metadata_list", &PrepareArgs::seq_metadata_list)
        .def_readwrite("block_size", &PrepareArgs::block_size)
        .def_readwrite("max_ngram_size", &PrepareArgs::max_ngram_size)
        .def_readwrite("engram_compressed_pad_token", &PrepareArgs::engram_compressed_pad_token)
        .def_readwrite("image_sentinel", &PrepareArgs::image_sentinel)
        .def_readwrite("is_prompt", &PrepareArgs::is_prompt)
        .def_readwrite("cp_size", &PrepareArgs::cp_size)
        .def_readwrite("max_num_blocks", &PrepareArgs::max_num_blocks)
        .def_readwrite("out", &PrepareArgs::out)
        .def_readwrite("window_size", &PrepareArgs::window_size)
        .def_readwrite("swa_block_size", &PrepareArgs::swa_block_size)
        .def_readwrite("yoco_prefill_truncation_mode", &PrepareArgs::yoco_prefill_truncation_mode)
        .def_readwrite("yoco_prefill_keep_tokens", &PrepareArgs::yoco_prefill_keep_tokens)
        .def_readwrite("max_backtracking_tokens", &PrepareArgs::max_backtracking_tokens)
        .def_readwrite("force_load_balance", &PrepareArgs::force_load_balance)
        .def_readwrite("dense_compress_ratios", &PrepareArgs::dense_compress_ratios)
        .def_readwrite("sparse_compress_ratios", &PrepareArgs::sparse_compress_ratios);

    engine.def("prepare_prompt_v4", &prepare_prompt_v4, py::arg("args"));
    engine.def("prepare_decode_v4", &prepare_decode_v4, py::arg("args"));
    engine.def("set_engram_lookup_table", &set_engram_lookup_table, py::arg("lookup_table"));
    py::class_<LogitsTransformInputs>(engine, "LogitsTransformInputs")
        .def(py::init([](torch::Tensor temperatures, torch::Tensor top_ps, torch::Tensor top_ks) {
            return LogitsTransformInputs{temperatures, top_ps, top_ks};
        }),
            py::arg("temperatures"), py::arg("top_ps"), py::arg("top_ks"))
        .def_readonly("temperatures", &LogitsTransformInputs::temperatures)
        .def_readonly("top_ps", &LogitsTransformInputs::top_ps)
        .def_readonly("top_ks", &LogitsTransformInputs::top_ks);
    py::class_<SamplingKernelInputs>(engine, "SamplingKernelInputs")
        .def(py::init([](LogitsTransformInputs logits, torch::Tensor seq_seeds, torch::Tensor sample_positions) {
            return SamplingKernelInputs{logits, seq_seeds, sample_positions};
        }),
            py::arg("logits"), py::arg("seq_seeds"), py::arg("sample_positions"))
        .def_readonly("logits", &SamplingKernelInputs::logits)
        .def_readonly("seq_seeds", &SamplingKernelInputs::seq_seeds)
        .def_readonly("sample_positions", &SamplingKernelInputs::sample_positions);
    py::class_<PromptLogprobInputs>(engine, "PromptLogprobInputs")
        .def(py::init(
                 [](torch::Tensor prompt_tokens, std::vector<int> prompt_logprobs_lens, std::vector<bool> enabled,
                     LogitsTransformInputs logits) { return PromptLogprobInputs{prompt_tokens, prompt_logprobs_lens, enabled, logits}; }),
            py::arg("prompt_tokens"), py::arg("prompt_logprobs_lens"), py::arg("enabled"), py::arg("logits"))
        .def_readonly("prompt_tokens", &PromptLogprobInputs::prompt_tokens)
        .def_readonly("prompt_logprobs_lens", &PromptLogprobInputs::prompt_logprobs_lens)
        .def_readonly("enabled", &PromptLogprobInputs::enabled)
        .def_readonly("logits", &PromptLogprobInputs::logits);
    py::class_<PrefillSamplingInputs>(engine, "PrefillSamplingInputs")
        .def_readonly("prompt_logprobs", &PrefillSamplingInputs::prompt_logprobs)
        .def_readonly("eos_suppression_rows", &PrefillSamplingInputs::eos_suppression_rows)
        .def_readonly("eos_token_ids", &PrefillSamplingInputs::eos_token_ids);
    py::class_<SpeculativeSamplingInputs>(engine, "SpeculativeSamplingInputs")
        .def_readonly("nextn_tokens", &SpeculativeSamplingInputs::nextn_tokens)
        .def_readonly("mtp_info_read_rows", &SpeculativeSamplingInputs::mtp_info_read_rows)
        .def_readonly("mtp_info_write_slots", &SpeculativeSamplingInputs::mtp_info_write_slots)
        .def_readonly("stop_token_ids", &SpeculativeSamplingInputs::stop_token_ids);
    py::class_<DecodeSamplingInputs>(engine, "DecodeSamplingInputs")
        .def(py::init([](torch::Tensor mtp_first_token_positions, torch::Tensor mtp_seqlens_q) {
            return DecodeSamplingInputs{mtp_first_token_positions, mtp_seqlens_q};
        }),
            py::arg("mtp_first_token_positions"), py::arg("mtp_seqlens_q"))
        .def_readonly("mtp_first_token_positions", &DecodeSamplingInputs::mtp_first_token_positions)
        .def_readonly("mtp_seqlens_q", &DecodeSamplingInputs::mtp_seqlens_q);
    py::class_<DriverSamplingInputs>(engine, "DriverSamplingInputs")
        .def_readonly("kernel", &DriverSamplingInputs::kernel)
        .def_readonly("max_top_logprobs", &DriverSamplingInputs::max_top_logprobs)
        .def_readonly("golden_output", &DriverSamplingInputs::golden_output)
        .def_readonly("prefill", &DriverSamplingInputs::prefill)
        .def_readonly("speculative", &DriverSamplingInputs::speculative);
    py::class_<SamplerInputs>(engine, "SamplerInputs")
        .def(py::init([](torch::Tensor selected_token_indices, std::optional<torch::Tensor> selected_prompt_indices,
                          std::optional<DecodeSamplingInputs> decode) {
            return SamplerInputs{selected_token_indices, selected_prompt_indices, decode, std::nullopt};
        }),
            py::arg("selected_token_indices"), py::arg("selected_prompt_indices"), py::arg("decode"))
        .def_readonly("selected_token_indices", &SamplerInputs::selected_token_indices)
        .def_readonly("selected_prompt_indices", &SamplerInputs::selected_prompt_indices)
        .def_readonly("decode", &SamplerInputs::decode)
        .def_readonly("driver", &SamplerInputs::driver);
    engine.def("prepare_sampler_inputs", &prepare_sampler_inputs, py::arg("seq_metadata_list"), py::arg("max_mtp_tokens"),
        py::arg_v("yoco_prefill_truncation_mode", YocoPrefillTruncationMode::OFF, "YocoPrefillTruncationMode.OFF"),
        py::arg("yoco_prefill_keep_tokens") = 0);
}

} // namespace deepinfer::engine