#pragma once

#include <cstdint>
#include <string>
#include <vector>
#include <memory>

namespace diarizen {

/// Frame-level speaker segmentation result.
struct SegmentResult {
    int num_frames = 199;
    int num_classes = 11;
    /// Flattened log-probabilities: result[frame * num_classes + class_idx].
    std::vector<float> log_probs;
};

/// DiariZen speaker segmentation: CNN NPU frontend + CPU backend.
class DiarizenSegmenter {
public:
    /// @param cnn_model_path  Path to cnn_features.axmodel.
    /// @param backend_onnx_path  Path to backend.onnx.
    DiarizenSegmenter(const std::string& cnn_model_path,
                      const std::string& backend_onnx_path);

    ~DiarizenSegmenter();

    /// Run segmentation on 4s of 16kHz mono audio.
    /// @param audio  64000 float32 samples (1-D, 16kHz mono, pre-normalized).
    /// @return Frame-level log-probabilities.
    SegmentResult run(const float* audio, int num_samples);

    int num_frames() const { return 199; }
    int num_classes() const { return 11; }

private:
    struct Impl;
    std::unique_ptr<Impl> impl_;
};

}  // namespace diarizen