#include "whisper.h"
#include <nlohmann/json.hpp>

#include <algorithm>
#include <chrono>
#include <cmath>
#include <cstdint>
#include <cstring>
#include <iostream>
#include <memory>
#include <string>
#include <thread>
#include <unordered_map>
#include <vector>

#if defined(_WIN32)
#include <fcntl.h>
#include <io.h>
#endif

using json = nlohmann::json;
using clock_type = std::chrono::steady_clock;

namespace {
constexpr std::size_t sample_rate = 16000;
constexpr std::size_t max_samples = sample_rate * 6;
constexpr std::size_t min_decode_samples = sample_rate / 2;
constexpr auto decode_interval = std::chrono::milliseconds(600);
constexpr std::uint64_t finalize_silence_us = 750000;

struct whisper_deleter { void operator()(whisper_context * value) const { if (value) whisper_free(value); } };
using whisper_pointer = std::unique_ptr<whisper_context, whisper_deleter>;

struct track_state {
    std::vector<float> samples;
    std::uint64_t audio_start_us = 0;
    std::uint64_t audio_end_us = 0;
    std::uint64_t captured_at_ms = 0;
    std::uint64_t last_voice_us = 0;
    std::uint64_t last_decode_us = 0;
    bool voiced = false;
    bool new_utterance = true;
};

struct word_analysis {
    std::string text;
    double confidence_sum = 0;
    int confidence_pieces = 0;
    std::int64_t t0 = 0;
    std::int64_t t1 = 0;
};

struct transcription_result {
    std::string text;
    std::vector<word_analysis> words;
};

whisper_pointer context;
std::string model_id = "unknown";
std::string backend = "cpu";
std::unordered_map<std::string, track_state> tracks;

void emit(const json & value) {
    std::cout << value.dump() << '\n' << std::flush;
}

void diagnostic(const std::string & state, const std::string & message) {
    emit({{"type", "diagnostic"}, {"state", state}, {"message", message}});
}

bool read_exact(char * destination, std::size_t size) {
    std::cin.read(destination, static_cast<std::streamsize>(size));
    return static_cast<std::size_t>(std::cin.gcount()) == size;
}

bool read_packet(json & metadata, std::vector<std::uint8_t> & pcm) {
    std::uint8_t header[8];
    if (!read_exact(reinterpret_cast<char *>(header), sizeof(header))) return false;
    auto read_u32 = [](const std::uint8_t * value) {
        return static_cast<std::uint32_t>(value[0]) | (static_cast<std::uint32_t>(value[1]) << 8U) |
               (static_cast<std::uint32_t>(value[2]) << 16U) | (static_cast<std::uint32_t>(value[3]) << 24U);
    };
    const auto json_size = read_u32(header);
    const auto pcm_size = read_u32(header + 4);
    if (json_size == 0 || json_size > 64 * 1024 || pcm_size > 5 * 32000 || pcm_size % 2 != 0) {
        diagnostic("packet_rejected", "Worker packet exceeded its framing limits.");
        return false;
    }
    std::string body(json_size, '\0');
    pcm.resize(pcm_size);
    if (!read_exact(body.data(), body.size()) || (pcm_size > 0 && !read_exact(reinterpret_cast<char *>(pcm.data()), pcm.size()))) return false;
    metadata = json::parse(body, nullptr, false);
    if (metadata.is_discarded() || !metadata.is_object()) {
        diagnostic("packet_rejected", "Worker metadata was not valid JSON.");
        return true;
    }
    return true;
}

std::string model_path(const json & model) {
    for (const char * key : {"path", "file", "filename"}) {
        if (model.contains(key) && model[key].is_string()) return model[key].get<std::string>();
    }
    return {};
}

void load_model(const json & packet) {
    const auto model = packet.value("model", json::object());
    const auto path = model_path(model);
    if (path.empty()) {
        emit({{"type", "model_error"}, {"code", "MODEL_PATH_REQUIRED"}, {"message", "The model packet did not include a local path."}});
        return;
    }
    auto parameters = whisper_context_default_params();
    context.reset(whisper_init_from_file_with_params(path.c_str(), parameters));
    if (!context) {
        emit({{"type", "model_error"}, {"code", "MODEL_LOAD_FAILED"}, {"message", "whisper.cpp could not load the selected model."}});
        return;
    }
    model_id = model.value("id", "unknown");
#if defined(LUMI_BACKEND_CUDA)
    backend = "cuda";
#else
    backend = "cpu";
#endif
    emit({{"type", "model_loaded"}, {"model_id", model_id}, {"backend", backend}});
}

transcription_result transcribe(const std::vector<float> & samples, double & inference_ms) {
    if (!context) return {};
    auto parameters = whisper_full_default_params(WHISPER_SAMPLING_GREEDY);
    parameters.print_realtime = false;
    parameters.print_progress = false;
    parameters.print_timestamps = false;
    parameters.print_special = false;
    parameters.translate = false;
    parameters.no_context = true;
    parameters.single_segment = true;
    parameters.token_timestamps = true;
    parameters.split_on_word = true;
    parameters.language = "en";
    parameters.n_threads = std::max(1U, std::min(4U, std::thread::hardware_concurrency()));
    const auto started = clock_type::now();
    const auto result = whisper_full(context.get(), parameters, samples.data(), static_cast<int>(samples.size()));
    inference_ms = std::chrono::duration<double, std::milli>(clock_type::now() - started).count();
    if (result != 0) return {};
    transcription_result output;
    const int segments = whisper_full_n_segments(context.get());
    for (int segment = 0; segment < segments; ++segment) {
        output.text += whisper_full_get_segment_text(context.get(), segment);
        const int tokens = whisper_full_n_tokens(context.get(), segment);
        for (int index = 0; index < tokens; ++index) {
            const char * value = whisper_full_get_token_text(context.get(), segment, index);
            if (!value || !*value) continue;
            std::string raw(value);
            if (raw.starts_with("<|")) continue;
            const auto first = raw.find_first_not_of(" \t\r\n");
            if (first == std::string::npos) continue;
            const auto last = raw.find_last_not_of(" \t\r\n");
            const bool begins_word = first > 0 || output.words.empty();
            const auto data = whisper_full_get_token_data(context.get(), segment, index);
            if (begins_word) output.words.push_back({});
            auto & word = output.words.back();
            word.text += raw.substr(first, last - first + 1);
            word.confidence_sum += std::clamp(static_cast<double>(data.p), 0.0, 1.0);
            word.confidence_pieces += 1;
            if (word.confidence_pieces == 1) word.t0 = std::max<std::int64_t>(0, data.t0);
            word.t1 = std::max(word.t0, data.t1);
        }
    }
    const auto first = output.text.find_first_not_of(" \t\r\n");
    const auto last = output.text.find_last_not_of(" \t\r\n");
    output.text = first == std::string::npos ? std::string{} : output.text.substr(first, last - first + 1);
    return output;
}

void decode_track(const std::string & session_id, const std::string & track_id, track_state & track, bool final) {
    if (!context || track.samples.size() < min_decode_samples) return;
    double inference_ms = 0;
    const auto result = transcribe(track.samples, inference_ms);
    if (!result.text.empty()) {
        const bool incomplete = !final && std::isalnum(static_cast<unsigned char>(result.text.back())) != 0;
        const auto emitted_at_ms = static_cast<std::uint64_t>(std::chrono::duration_cast<std::chrono::milliseconds>(
            std::chrono::system_clock::now().time_since_epoch()).count());
        json words = json::array();
        for (const auto & word : result.words) {
            const auto captured_at_ms = track.captured_at_ms + static_cast<std::uint64_t>(word.t0) * 10ULL;
            words.push_back({
                {"text", word.text},
                {"confidence", word.confidence_pieces ? word.confidence_sum / word.confidence_pieces : 0.0},
                {"audio_start_ms", static_cast<double>(word.t0) * 10.0},
                {"audio_end_ms", static_cast<double>(word.t1) * 10.0},
                {"captured_at_ms", captured_at_ms},
                {"latency_ms", emitted_at_ms >= captured_at_ms ? emitted_at_ms - captured_at_ms : 0}
            });
        }
        emit({
            {"type", "hypothesis"}, {"session_id", session_id}, {"track_id", track_id}, {"text", result.text},
            {"final", final}, {"new_utterance", track.new_utterance}, {"incomplete_word", incomplete},
            {"audio_start_us", track.audio_start_us}, {"audio_end_us", track.audio_end_us},
            {"inference_ms", inference_ms}, {"model_id", model_id}, {"backend", backend},
            {"emitted_at_ms", emitted_at_ms}, {"words", std::move(words)}
        });
        track.new_utterance = false;
    }
    track.last_decode_us = track.audio_end_us;
    if (final) {
        track.samples.clear();
        track.voiced = false;
        track.new_utterance = true;
        track.audio_start_us = 0;
        track.captured_at_ms = 0;
    }
}

void audio(const json & packet, const std::vector<std::uint8_t> & pcm) {
    if (!context || pcm.empty()) return;
    const auto session_id = packet.value("session_id", "");
    const auto track_id = packet.value("track_id", "");
    if (session_id.empty() || track_id.empty()) return;
    auto & track = tracks[session_id + "\n" + track_id];
    const auto captured_us = packet.value("capture_timestamp_us", std::uint64_t{0});
    if (track.samples.empty()) {
        track.audio_start_us = captured_us;
        track.captured_at_ms = packet.value("captured_at", std::uint64_t{0});
    }
    const auto count = pcm.size() / 2;
    double squares = 0;
    for (std::size_t index = 0; index < count; ++index) {
        const auto sample = static_cast<std::int16_t>(static_cast<std::uint16_t>(pcm[index * 2]) |
            (static_cast<std::uint16_t>(pcm[index * 2 + 1]) << 8U));
        const float normalized = static_cast<float>(sample) / 32768.0F;
        track.samples.push_back(normalized);
        squares += static_cast<double>(normalized) * normalized;
    }
    const auto duration_us = static_cast<std::uint64_t>(count * 1000000ULL / sample_rate);
    track.audio_end_us = captured_us + duration_us;
    const auto rms = std::sqrt(squares / static_cast<double>(std::max<std::size_t>(1, count)));
    if (rms >= 0.012) { track.voiced = true; track.last_voice_us = track.audio_end_us; }
    const bool final = track.voiced && (track.samples.size() >= max_samples ||
        (track.audio_end_us > track.last_voice_us && track.audio_end_us - track.last_voice_us >= finalize_silence_us));
    const bool decode_due = track.last_decode_us == 0 || track.audio_end_us - track.last_decode_us >= static_cast<std::uint64_t>(decode_interval.count()) * 1000ULL;
    if (final || (track.voiced && decode_due)) decode_track(session_id, track_id, track, final);
}

void handle(const json & packet, const std::vector<std::uint8_t> & pcm, bool & running) {
    const auto type = packet.value("type", "");
    if (type == "load_model") load_model(packet);
    else if (type == "benchmark") emit({{"type", "benchmark"}, {"accepted", context != nullptr}, {"model_id", model_id}, {"backend", backend}});
    else if (type == "start_session") emit({{"type", "session_started"}, {"session_id", packet.value("session", json::object()).value("id", "")}});
    else if (type == "add_track") tracks.try_emplace(packet.value("session_id", "") + "\n" + packet.value("track", json::object()).value("source_uuid", ""));
    else if (type == "audio") audio(packet, pcm);
    else if (type == "remove_track") tracks.erase(packet.value("session_id", "") + "\n" + packet.value("track_id", ""));
    else if (type == "stop_session") {
        const auto session_id = packet.value("session_id", "");
        const auto prefix = session_id + "\n";
        for (auto & [key, track] : tracks) {
            if (key.starts_with(prefix) && track.voiced) decode_track(session_id, key.substr(prefix.size()), track, true);
        }
        std::erase_if(tracks, [&](const auto & item) { return item.first.starts_with(prefix); });
        emit({{"type", "session_stopped"}, {"session_id", session_id}});
    }
    else if (type == "shutdown") running = false;
    else diagnostic("packet_rejected", "Unknown worker message type.");
}
}

int main() {
#if defined(_WIN32)
    // PCM frames are transported over stdin. Windows text mode treats 0x1A in
    // ordinary audio as EOF, which closes the worker and breaks the host pipe.
    if (_setmode(_fileno(stdin), _O_BINARY) == -1) return 2;
#endif
    std::ios::sync_with_stdio(false);
    whisper_log_set([](enum ggml_log_level, const char *, void *) {}, nullptr);
    emit({{"type", "ready"}, {"protocol_version", 1}, {"backend", backend}});
    bool running = true;
    while (running && std::cin.good()) {
        json packet;
        std::vector<std::uint8_t> pcm;
        if (!read_packet(packet, pcm)) break;
        if (!packet.is_discarded()) handle(packet, pcm, running);
    }
    tracks.clear();
    context.reset();
    return 0;
}
