#pragma once

#include "engine/framework/assets/tensor_source.h"
#include "engine/framework/core/backend.h"
#include "engine/framework/core/module.h"
#include "engine/framework/modules/conv_modules.h"
#include "engine/framework/runtime/model.h"
#include "engine/models/dramabox/assets.h"

#include <memory>
#include <optional>
#include <vector>

namespace engine::core {
class BackendWeightStore;
class ExecutionContext;
}

namespace engine::models::dramabox {

struct DramaBoxVAEResnetBlockWeights {
    modules::PixelNormCausalConv2dResBlockConfig config;
    modules::PixelNormCausalConv2dResBlockWeights block;
};

struct DramaBoxVAEUpStageWeights {
    std::vector<DramaBoxVAEResnetBlockWeights> blocks;
    std::optional<modules::Conv2dWeights> upsample;
    int64_t channels = 0;
};

struct DramaBoxAudioVAEDecoderWeights {
    std::shared_ptr<core::BackendWeightStore> store;
    std::vector<float> latent_mean;
    std::vector<float> latent_std;
    modules::Conv2dWeights conv_in;
    DramaBoxVAEResnetBlockWeights mid_block_1;
    DramaBoxVAEResnetBlockWeights mid_block_2;
    std::vector<DramaBoxVAEUpStageWeights> up;
    modules::Conv2dWeights conv_out;
};

struct DramaBoxVAEDownStageWeights {
    std::vector<DramaBoxVAEResnetBlockWeights> blocks;
    std::optional<modules::Conv2dWeights> downsample;
    int64_t channels = 0;
};

struct DramaBoxAudioVAEEncoderWeights {
    std::shared_ptr<core::BackendWeightStore> store;
    std::vector<float> latent_mean;
    std::vector<float> latent_std;
    modules::Conv2dWeights conv_in;
    std::vector<DramaBoxVAEDownStageWeights> down;
    DramaBoxVAEResnetBlockWeights mid_block_1;
    DramaBoxVAEResnetBlockWeights mid_block_2;
    modules::Conv2dWeights conv_out;
};

struct DramaBoxDecodedMel {
    int64_t batch = 0;
    int64_t channels = 0;
    int64_t frames = 0;
    int64_t mel_bins = 0;
    std::vector<float> values;
    const ggml_tensor * device_values = nullptr;
};

struct DramaBoxEncodedReferenceLatents {
    int64_t tokens = 0;
    std::vector<float> values;
};

DramaBoxAudioVAEDecoderWeights load_dramabox_audio_vae_decoder_weights(
    const DramaBoxAssets & assets,
    ggml_backend_t backend,
    core::BackendType backend_type,
    size_t weight_context_bytes,
    assets::TensorStorageType weight_storage_type);

DramaBoxAudioVAEEncoderWeights load_dramabox_audio_vae_encoder_weights(
    const DramaBoxAssets & assets,
    ggml_backend_t backend,
    core::BackendType backend_type,
    size_t weight_context_bytes,
    assets::TensorStorageType weight_storage_type);

std::vector<float> reference_log_mel(
    const runtime::AudioBuffer & audio_buffer,
    const DramaBoxConfig & config,
    float ref_duration,
    int threads,
    int64_t & frames_out);

class DramaBoxAudioVAEDecoderRuntime {
public:
    DramaBoxAudioVAEDecoderRuntime(
        core::ExecutionContext & execution,
        std::shared_ptr<const DramaBoxAssets> assets,
        assets::TensorStorageType weight_storage_type);
    ~DramaBoxAudioVAEDecoderRuntime();

    DramaBoxAudioVAEDecoderRuntime(const DramaBoxAudioVAEDecoderRuntime &) = delete;
    DramaBoxAudioVAEDecoderRuntime & operator=(const DramaBoxAudioVAEDecoderRuntime &) = delete;

    void prepare(int64_t batch, int64_t latent_frames) const;
    DramaBoxDecodedMel decode(const std::vector<float> & patch_latents, int64_t batch, int64_t latent_frames) const;
    DramaBoxDecodedMel decode_to_device(const std::vector<float> & patch_latents, int64_t batch, int64_t latent_frames) const;
    void release_runtime_state() const;

private:
    class Graph;

    core::ExecutionContext * execution_ = nullptr;
    std::shared_ptr<const DramaBoxAssets> assets_;
    assets::TensorStorageType weight_storage_type_ = assets::TensorStorageType::Native;
    mutable std::unique_ptr<DramaBoxAudioVAEDecoderWeights> weights_;
    mutable std::unique_ptr<Graph> graph_;
};

class DramaBoxAudioVAEEncoderRuntime {
public:
    DramaBoxAudioVAEEncoderRuntime(
        core::ExecutionContext & execution,
        std::shared_ptr<const DramaBoxAssets> assets,
        assets::TensorStorageType weight_storage_type);
    ~DramaBoxAudioVAEEncoderRuntime();

    DramaBoxAudioVAEEncoderRuntime(const DramaBoxAudioVAEEncoderRuntime &) = delete;
    DramaBoxAudioVAEEncoderRuntime & operator=(const DramaBoxAudioVAEEncoderRuntime &) = delete;

    void prepare(int64_t batch, int64_t mel_frames) const;
    DramaBoxEncodedReferenceLatents encode(const std::vector<float> & mel, int64_t batch, int64_t mel_frames) const;
    void release_runtime_state() const;

private:
    class Graph;

    core::ExecutionContext * execution_ = nullptr;
    std::shared_ptr<const DramaBoxAssets> assets_;
    assets::TensorStorageType weight_storage_type_ = assets::TensorStorageType::Native;
    mutable std::unique_ptr<DramaBoxAudioVAEEncoderWeights> weights_;
    mutable std::unique_ptr<Graph> graph_;
};

}  // namespace engine::models::dramabox