#pragma once

#include "engine/framework/assets/tensor_source.h"

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

namespace engine::models::marblenet_vad {

struct JasperBlockConfig {
    int64_t filters = 0;
    int64_t repeat = 0;
    int64_t kernel = 0;
    int64_t stride = 1;
    int64_t dilation = 1;
    bool residual = false;
    bool separable = false;
};

struct MarbleNetConfig {
    int64_t sample_rate = 0;
    int64_t n_mels = 0;
    int64_t n_fft = 0;
    int64_t hop_length = 0;
    int64_t win_length = 0;
    int64_t pad_to = 0;
    int64_t num_classes = 0;
    int64_t output_stride = 1;
    std::string normalize;
    std::string window;
    std::vector<JasperBlockConfig> jasper;
};

struct BatchNorm1dWeights {
    int64_t channels = 0;
    std::vector<float> weight;
    std::vector<float> bias;
    std::vector<float> running_mean;
    std::vector<float> running_var;
};

struct Conv1dWeights {
    std::string weight_name;
    std::vector<int64_t> weight_source_shape;
    int64_t in_channels = 0;
    int64_t out_channels = 0;
    int64_t kernel = 0;
    int64_t stride = 1;
    int64_t dilation = 1;
    int64_t padding = 0;
    int64_t groups = 1;
    bool use_bias = false;
    std::optional<std::string> bias_name;
};

struct ConvBn {
    Conv1dWeights conv;
    BatchNorm1dWeights bn;
};

struct SeparableConvBn {
    Conv1dWeights depthwise;
    Conv1dWeights pointwise;
    BatchNorm1dWeights bn;
};

struct JasperBlockWeights {
    bool separable = false;
    std::vector<SeparableConvBn> separable_repeats;
    std::vector<ConvBn> conv_repeats;
    bool has_residual = false;
    Conv1dWeights residual_conv;
    BatchNorm1dWeights residual_bn;
};

struct MarbleNetWeights {
    MarbleNetConfig config;
    std::shared_ptr<const assets::TensorSource> source;
    std::vector<float> window;
    std::vector<float> fb;
    std::vector<JasperBlockWeights> blocks;
    Conv1dWeights decoder;
};

struct MarbleNetAssetPaths {
    std::filesystem::path model_root;
    std::filesystem::path checkpoint_path;
    std::filesystem::path config_path;
    std::filesystem::path labels_path;
};

MarbleNetAssetPaths resolve_marblenet_assets(const std::filesystem::path & checkpoint_path);
std::shared_ptr<const MarbleNetWeights> load_marblenet_weights_cached(const std::filesystem::path & checkpoint_path);

}  // namespace engine::models::marblenet_vad