#pragma once

#include <array>
#include <cstddef>
#include <cstdint>

namespace ling3 {

enum class AttentionKind : std::uint8_t { kKda, kMla };
enum class FfnKind : std::uint8_t { kDense, kMoe };

struct LayerPlan {
    std::uint8_t index = 0;
    AttentionKind attention = AttentionKind::kKda;
    FfnKind ffn = FfnKind::kMoe;
};

struct ExecutionPlan {
    std::array<LayerPlan, 24> layers{};
    std::size_t decode_matmul_contexts = 0;
    std::size_t rknn_island_contexts = 0;
    std::size_t kda_state_bytes_fp16 = 0;
    std::size_t mla_cache_bytes_fp16 = 0;
    std::size_t activation_arena_bytes = 0;
};

ExecutionPlan BuildExecutionPlan(std::size_t max_context);

} // namespace ling3