#include "engine/models/higgs_audio_stt/session.h"

#include "engine/framework/audio/chunking.h"
#include "engine/framework/debug/profiler.h"
#include "engine/framework/runtime/options.h"
#include "engine/framework/runtime/spec_backed_model.h"

#include <algorithm>
#include <chrono>
#include <cmath>
#include <cstdint>
#include <sstream>
#include <stdexcept>
#include <string>
#include <unordered_map>
#include <utility>

namespace engine::models::higgs_audio_stt {
namespace {

using Clock = std::chrono::steady_clock;

constexpr const char * kFamily = "higgs_audio_stt";
constexpr size_t kDefaultAudioEncoderGraphArenaBytes = 512ull * 1024ull * 1024ull;
constexpr size_t kDefaultTextDecoderPrefillGraphArenaBytes = 512ull * 1024ull * 1024ull;
constexpr size_t kDefaultTextDecoderDecodeGraphArenaBytes = 256ull * 1024ull * 1024ull;
#if defined(INTPTR_MAX) && (INTPTR_MAX == INT32_MAX)
constexpr size_t kDefaultTextDecoderWeightContextBytes = 1024ull * 1024ull * 1024ull;
#else
constexpr size_t kDefaultTextDecoderWeightContextBytes = 4096ull * 1024ull * 1024ull;
#endif

std::shared_ptr<const HiggsAudioSTTAssets> require_assets(std::shared_ptr<const HiggsAudioSTTAssets> assets) {
    if (assets == nullptr) {
        throw std::runtime_error("Higgs Audio STT session requires assets");
    }
    return assets;
}

std::shared_ptr<const engine::model_spec::ModelContract> require_contract(
    std::shared_ptr<const engine::model_spec::ModelContract> contract) {
    if (contract == nullptr) {
        throw std::runtime_error("Higgs Audio STT session requires a model contract");
    }
    return contract;
}

int64_t audio_frame_count(const runtime::AudioBuffer & audio) {
    if (audio.channels <= 0) {
        throw std::runtime_error("Higgs Audio STT audio chunking requires positive audio channels");
    }
    if (audio.samples.size() % static_cast<size_t>(audio.channels) != 0) {
        throw std::runtime_error("Higgs Audio STT audio samples must be divisible by channel count");
    }
    return static_cast<int64_t>(audio.samples.size() / static_cast<size_t>(audio.channels));
}

runtime::SessionOptions normalize_session_options(
    runtime::SessionOptions options,
    const std::shared_ptr<const engine::model_spec::ModelContract> & contract) {
    options = runtime::apply_option_v1_compatibility(
        std::move(options),
        {
            {"weight_type", "higgs_audio_stt.weight_type"},
            {"audio_encoder_weight_type", "higgs_audio_stt.audio_encoder_weight_type"},
            {"text_decoder_weight_type", "higgs_audio_stt.text_decoder_weight_type"},
            {"audio_encoder_graph_arena_mb", "higgs_audio_stt.audio_encoder_graph_arena_mb"},
            {"text_decoder_prefill_graph_arena_mb", "higgs_audio_stt.text_decoder_prefill_graph_arena_mb"},
            {"text_decoder_decode_graph_arena_mb", "higgs_audio_stt.text_decoder_decode_graph_arena_mb"},
            {"text_decoder_weight_context_mb", "higgs_audio_stt.text_decoder_weight_context_mb"},
        },
        "Higgs Audio STT");
    runtime::validate_spec_backed_session_options(
        options,
        *require_contract(contract),
        kFamily,
        "Higgs Audio STT");
    return options;
}

std::unordered_map<std::string, std::string> normalize_request_options(
    std::unordered_map<std::string, std::string> options) {
    return runtime::apply_option_v1_compatibility(
        std::move(options),
        {
            {"audio_chunk_seconds", "audio_chunk_duration_sec"},
            {"audio_chunk_duration_seconds", "audio_chunk_duration_sec"},
            {"audio_chunk_duration", "audio_chunk_duration_sec"},
        },
        "Higgs Audio STT",
        "request");
}

std::unique_ptr<runtime::IVoiceTaskSession> create_higgs_audio_stt_session(
    const runtime::TaskSpec & task,
    const runtime::SessionOptions & options,
    std::shared_ptr<const HiggsAudioSTTAssets> assets,
    std::shared_ptr<const engine::model_spec::ModelContract> contract) {
    return std::make_unique<HiggsAudioSTTSession>(
        task,
        options,
        std::move(assets),
        std::move(contract));
}

}  // namespace

HiggsAudioSTTSession::HiggsAudioSTTSession(
    runtime::TaskSpec task,
    runtime::SessionOptions options,
    std::shared_ptr<const HiggsAudioSTTAssets> assets,
    std::shared_ptr<const engine::model_spec::ModelContract> contract)
    : RuntimeSessionBase(normalize_session_options(std::move(options), contract)),
      task_(task),
      assets_(require_assets(std::move(assets))),
      contract_(require_contract(std::move(contract))),
      audio_encoder_graph_arena_bytes_(runtime::parse_size_mb_option(RuntimeSessionBase::options().options, {"higgs_audio_stt.audio_encoder_graph_arena_mb"}, kDefaultAudioEncoderGraphArenaBytes)),
      text_decoder_prefill_graph_arena_bytes_(runtime::parse_size_mb_option(RuntimeSessionBase::options().options, {"higgs_audio_stt.text_decoder_prefill_graph_arena_mb"}, kDefaultTextDecoderPrefillGraphArenaBytes)),
      text_decoder_decode_graph_arena_bytes_(runtime::parse_size_mb_option(RuntimeSessionBase::options().options, {"higgs_audio_stt.text_decoder_decode_graph_arena_mb"}, kDefaultTextDecoderDecodeGraphArenaBytes)),
      text_decoder_weight_context_bytes_(runtime::parse_size_mb_option(RuntimeSessionBase::options().options, {"higgs_audio_stt.text_decoder_weight_context_mb"}, kDefaultTextDecoderWeightContextBytes)),
      audio_encoder_weight_storage_type_(runtime::parse_tensor_storage_option(
          RuntimeSessionBase::options().options,
          "higgs_audio_stt.audio_encoder_weight_type",
          engine::assets::TensorStorageType::Native,
          {
              engine::assets::TensorStorageType::Native,
              engine::assets::TensorStorageType::F32,
              engine::assets::TensorStorageType::F16,
          })),
      text_decoder_weight_storage_type_(runtime::parse_tensor_storage_option(
          RuntimeSessionBase::options().options,
          "higgs_audio_stt.text_decoder_weight_type",
          "higgs_audio_stt.weight_type",
          engine::assets::TensorStorageType::Native,
          {
              engine::assets::TensorStorageType::Native,
              engine::assets::TensorStorageType::F32,
              engine::assets::TensorStorageType::F16,
              engine::assets::TensorStorageType::BF16,
              engine::assets::TensorStorageType::Q8_0,
          })),
      tokenizer_(assets_),
      frontend_(assets_),
      audio_encoder_(assets_, execution_context(), audio_encoder_graph_arena_bytes_, audio_encoder_weight_storage_type_),
      text_decoder_(
          assets_,
          execution_context(),
          text_decoder_prefill_graph_arena_bytes_,
          text_decoder_decode_graph_arena_bytes_,
          text_decoder_weight_context_bytes_,
          text_decoder_weight_storage_type_),
      prompt_builder_(tokenizer_),
      postprocessor_(tokenizer_) {
    if (task_.task != runtime::VoiceTaskKind::Asr) {
        throw std::runtime_error("Higgs Audio STT only supports VoiceTaskKind::Asr");
    }
    if (task_.mode != runtime::RunMode::Offline && task_.mode != runtime::RunMode::Streaming) {
        throw std::runtime_error("Higgs Audio STT supports offline and streaming sessions");
    }
    assets_->model_weights->release_storage();
}

HiggsAudioSTTSession::~HiggsAudioSTTSession() = default;

std::string HiggsAudioSTTSession::family() const {
    return kFamily;
}

runtime::VoiceTaskKind HiggsAudioSTTSession::task_kind() const {
    return task_.task;
}

runtime::RunMode HiggsAudioSTTSession::run_mode() const {
    return task_.mode;
}

void HiggsAudioSTTSession::prepare(const runtime::SessionPreparationRequest & request) {
    const auto prepare_start = Clock::now();
    if (!request.audio.has_value()) {
        throw std::runtime_error("Higgs Audio STT prepare() requires an audio contract");
    }
    mark_prepared();
    debug::timing_log_scalar("higgs_audio_stt.prepare_ms", engine::debug::elapsed_ms(prepare_start, Clock::now()));
    debug::trace_log_scalar("higgs_audio_stt.prepare.max_input_samples", request.audio->max_input_samples);
}

runtime::TaskResult HiggsAudioSTTSession::run(const runtime::TaskRequest & request) {
    require_prepared("Higgs Audio STT run()");
    auto normalized_request = request;
    normalized_request.options = normalize_request_options(request.options);
    runtime::validate_spec_backed_request_options(
        normalized_request.options,
        *contract_,
        "Higgs Audio STT");
    if (task_.mode != runtime::RunMode::Offline) {
        throw std::runtime_error("Higgs Audio STT offline run called on non-offline session");
    }
    const auto chunks = audio_chunk_plan(normalized_request);
    if (chunks.empty()) {
        return run_single(make_request(normalized_request));
    }
    const auto & audio = *normalized_request.audio_input;
    if (chunks.size() == 1) {
        auto item_request = normalized_request;
        item_request.audio_input = engine::audio::slice_audio_buffer(audio, chunks.front().source_span);
        return run_single(make_request(item_request));
    }
    runtime::TaskResult merged;
    std::ostringstream text;
    for (const auto & chunk : chunks) {
        auto item_request = normalized_request;
        item_request.audio_input = engine::audio::slice_audio_buffer(audio, chunk.source_span);
        const auto item = run_single(make_request(item_request));
        if (item.text_output.has_value() && !item.text_output->text.empty()) {
            if (text.tellp() > 0) {
                text << ' ';
            }
            text << item.text_output->text;
            if (!merged.text_output.has_value()) {
                merged.text_output = runtime::Transcript{"", item.text_output->language};
            } else if (merged.text_output->language.empty()) {
                merged.text_output->language = item.text_output->language;
            }
        }
    }
    if (merged.text_output.has_value()) {
        merged.text_output->text = text.str();
    }
    return merged;
}

runtime::StreamingPolicy HiggsAudioSTTSession::streaming_policy() const {
    runtime::StreamingPolicy policy;
    policy.input = runtime::StreamingInputKind::AudioChunks;
    policy.output = runtime::StreamingOutputKind::FinalResult;
    policy.preferred_audio_chunk_seconds = 4.0;
    return policy;
}

void HiggsAudioSTTSession::begin_stream(const runtime::TaskRequest & request) {
    require_prepared("Higgs Audio STT start_stream()");
    streaming_request_ = request;
    streaming_request_.options = normalize_request_options(request.options);
    runtime::validate_spec_backed_request_options(
        streaming_request_.options,
        *contract_,
        "Higgs Audio STT");
    streaming_request_.audio_input = std::nullopt;
    streaming_chunks_processed_ = 0;
}

void HiggsAudioSTTSession::reset_stream_state() {
    require_prepared("Higgs Audio STT reset()");
    if (task_.mode != runtime::RunMode::Streaming) {
        throw std::runtime_error("Higgs Audio STT reset called on non-streaming session");
    }
    streaming_request_ = runtime::TaskRequest{};
    streaming_chunks_processed_ = 0;
}

void HiggsAudioSTTSession::consume_audio(const runtime::AudioChunk & chunk) {
    require_prepared("Higgs Audio STT process_audio_chunk()");
    runtime::AudioBuffer audio;
    audio.sample_rate = chunk.sample_rate;
    audio.channels = chunk.channels;
    audio.samples = chunk.samples;

    auto request = streaming_request_;
    request.audio_input = std::move(audio);
    auto result = run_single(make_request(request));
    ++streaming_chunks_processed_;

    if (result.text_output.has_value()) {
        auto & transcript = *result.text_output;
        if (!transcript.text.empty() && !committed_transcript().text.empty()) {
            transcript.text.insert(0, " ");
        }
        append_transcript(std::move(transcript));
    }
}

void HiggsAudioSTTSession::complete_stream(runtime::TaskResult &) {
    require_prepared("Higgs Audio STT finalize()");
    if (streaming_chunks_processed_ == 0) {
        throw std::runtime_error("Higgs Audio STT finalize requires streamed audio");
    }
}

runtime::TaskResult HiggsAudioSTTSession::run_single(const HiggsAudioSTTRequest & asr_request) {
    const auto wall_start = Clock::now();
    const auto frontend_start = Clock::now();
    const auto features = frontend_.extract(asr_request.audio);
    const auto frontend_end = Clock::now();
    const auto prompt_start = Clock::now();
    const auto prompt = prompt_builder_.build(asr_request, features.encoder_tokens);
    const auto prompt_end = Clock::now();
    const auto encoder_start = Clock::now();
    const auto audio_embeddings = audio_encoder_.encode(features);
    const auto encoder_end = Clock::now();
    const auto text_decoder_start = Clock::now();
    const auto tokens = text_decoder_.generate(prompt, audio_embeddings, asr_request.generation);
    const auto text_decoder_end = Clock::now();
    const auto postprocess_start = Clock::now();
    const auto decoded = postprocessor_.decode(tokens, asr_request);
    const auto postprocess_end = Clock::now();

    runtime::TaskResult result;
    result.text_output = runtime::Transcript{decoded.text, decoded.language};

    debug::timing_log_scalar("higgs_audio_stt.frontend_ms", engine::debug::elapsed_ms(frontend_start, frontend_end));
    debug::timing_log_scalar("higgs_audio_stt.prompt_ms", engine::debug::elapsed_ms(prompt_start, prompt_end));
    debug::timing_log_scalar("higgs_audio_stt.audio_encoder_ms", engine::debug::elapsed_ms(encoder_start, encoder_end));
    debug::timing_log_scalar("higgs_audio_stt.text_decoder_ms", engine::debug::elapsed_ms(text_decoder_start, text_decoder_end));
    debug::timing_log_scalar("higgs_audio_stt.postprocess_ms", engine::debug::elapsed_ms(postprocess_start, postprocess_end));
    debug::timing_log_scalar("session.wall_ms", engine::debug::elapsed_ms(wall_start, Clock::now()));
    debug::trace_log_scalar("higgs_audio_stt.prompt_tokens", prompt.input_ids.size());
    debug::trace_log_scalar("higgs_audio_stt.audio_chunks", features.chunks);
    debug::trace_log_scalar("higgs_audio_stt.audio_frames", features.frames);
    debug::trace_log_scalar("higgs_audio_stt.audio_tokens", features.encoder_tokens);
    return result;
}

std::vector<HiggsAudioSTTSession::AudioChunkPlan> HiggsAudioSTTSession::audio_chunk_plan(
    const runtime::TaskRequest & request) const {
    if (!request.audio_input.has_value()) {
        return {};
    }
    const auto mode = engine::audio::parse_audio_chunk_mode(request.options);
    if (mode == engine::audio::AudioChunkMode::None) {
        return {};
    }
    if (mode == engine::audio::AudioChunkMode::QuietEnergy ||
        mode == engine::audio::AudioChunkMode::Vad) {
        throw std::runtime_error("Higgs Audio STT supports audio_chunk_mode=auto, fixed, or none");
    }
    const auto & audio = *request.audio_input;
    const int64_t frames = audio_frame_count(audio);
    const auto seconds = engine::audio::parse_audio_chunk_seconds_override(request.options).value_or(4.0F);
    if (!(seconds > 0.0F)) {
        throw std::runtime_error("Higgs Audio STT audio_chunk_duration_sec must be positive");
    }
    const int64_t samples = static_cast<int64_t>(
        std::llround(static_cast<double>(seconds) * static_cast<double>(audio.sample_rate)));
    if (samples <= 0) {
        throw std::runtime_error("Higgs Audio STT audio_chunk_duration_sec produced an empty chunk");
    }
    const auto chunks = engine::audio::plan_audio_chunks(
        frames,
        engine::audio::AudioChunkSpec{
            samples,
            samples,
            engine::audio::AudioChunkPadMode::Zero,
            engine::audio::AudioChunkTailAlignment::Start,
            0,
        });
    std::vector<AudioChunkPlan> plan;
    plan.reserve(chunks.size());
    for (const auto & chunk : chunks) {
        plan.push_back(AudioChunkPlan{
            runtime::TimeSpan{
                chunk.output_start_sample,
                chunk.output_start_sample + chunk.valid_samples,
            },
        });
    }
    return plan;
}

HiggsAudioSTTRequest HiggsAudioSTTSession::make_request(const runtime::TaskRequest & request) const {
    if (!request.audio_input.has_value()) {
        throw std::runtime_error("Higgs Audio STT run() requires audio_input");
    }
    HiggsAudioSTTRequest out;
    out.audio = *request.audio_input;
    out.generation.max_new_tokens = assets_->config.max_new_tokens;
    if (request.text_input.has_value()) {
        out.context = request.text_input->text;
        out.language = request.text_input->language;
    }
    if (const auto language = runtime::find_option(request.options, {"language"})) {
        out.language = *language;
    }
    if (const auto value = runtime::parse_int_option(request.options, {"max_tokens"})) {
        out.generation.max_new_tokens = *value;
        if (out.generation.max_new_tokens <= 0) {
            throw std::runtime_error("Higgs Audio STT max_tokens must be positive");
        }
    }
    if (const auto value = runtime::find_option(request.options, {"enable_thinking"})) {
        out.generation.enable_thinking = runtime::parse_bool_option(*value, "enable_thinking");
    }
    return out;
}

// Loading adapter: Higgs Audio STT uses the schema-v1 spec-backed loader, so the
// loader wiring stays beside the session it constructs.
std::shared_ptr<runtime::IVoiceModelLoader> make_higgs_audio_stt_loader() {
    runtime::SpecBackedVoiceModelConfig<HiggsAudioSTTAssets> config;
    config.family = kFamily;
    config.load_assets = load_higgs_audio_stt_assets;
    config.create_session = create_higgs_audio_stt_session;
    return runtime::make_spec_backed_voice_loader(std::move(config));
}

}  // namespace engine::models::higgs_audio_stt