#include "VoskSTTEngine.hpp"

// Check if Vosk is available at compile time
#ifdef HAS_VOSK
#include <vosk_api.h>
#endif

#include <nlohmann/json.hpp>
#include <spdlog/spdlog.h>
#include <spdlog/sinks/stdout_color_sinks.h>
#include <cstring>
#include <algorithm>

namespace celuna {

VoskSTTEngine::VoskSTTEngine(const std::string& modelPath, float sampleRate)
    : m_sampleRate(sampleRate)
{
    m_logger = spdlog::get("VoskSTT");
    if (!m_logger) {
        m_logger = spdlog::stdout_color_mt("VoskSTT");
    }

#ifdef HAS_VOSK
    m_logger->info("Initializing Vosk STT engine...");
    m_logger->info("Model path: {}", modelPath);
    m_logger->info("Sample rate: {} Hz", sampleRate);

    // Load Vosk model
    m_model = vosk_model_new(modelPath.c_str());
    if (!m_model) {
        m_logger->error("Failed to load Vosk model from: {}", modelPath);
        m_logger->error("Make sure the model is downloaded and path is correct");
        m_logger->error("Download: https://alphacephei.com/vosk/models");
        m_available = false;
        return;
    }

    // Create recognizer
    m_recognizer = vosk_recognizer_new(m_model, sampleRate);
    if (!m_recognizer) {
        m_logger->error("Failed to create Vosk recognizer");
        vosk_model_free(m_model);
        m_model = nullptr;
        m_available = false;
        return;
    }

    m_available = true;
    m_logger->info("Vosk STT engine initialized successfully");
#else
    m_logger->warn("Vosk support not compiled (HAS_VOSK not defined)");
    m_logger->warn("To enable Vosk: install libvosk-dev and rebuild with -DHAS_VOSK");
    m_available = false;
#endif
}

VoskSTTEngine::~VoskSTTEngine() {
#ifdef HAS_VOSK
    if (m_recognizer) {
        vosk_recognizer_free(m_recognizer);
        m_recognizer = nullptr;
    }
    if (m_model) {
        vosk_model_free(m_model);
        m_model = nullptr;
    }
    m_logger->debug("Vosk STT engine destroyed");
#endif
}

std::string VoskSTTEngine::transcribe(const std::vector<float>& audioData) {
#ifdef HAS_VOSK
    if (!m_available || audioData.empty()) {
        return "";
    }

    // Convert float [-1.0, 1.0] to int16 [-32768, 32767]
    std::vector<int16_t> samples(audioData.size());
    for (size_t i = 0; i < audioData.size(); ++i) {
        float clamped = std::max(-1.0f, std::min(1.0f, audioData[i]));
        samples[i] = static_cast<int16_t>(clamped * 32767.0f);
    }

    // Feed audio to recognizer
    int accepted = vosk_recognizer_accept_waveform(
        m_recognizer,
        reinterpret_cast<const char*>(samples.data()),
        samples.size() * sizeof(int16_t)
    );

    // Get result
    const char* result = nullptr;
    if (accepted) {
        result = vosk_recognizer_result(m_recognizer);
    } else {
        result = vosk_recognizer_partial_result(m_recognizer);
    }

    if (!result) {
        return "";
    }

    // Parse JSON result
    std::string text = parseVoskResult(result);

    if (!text.empty()) {
        m_logger->debug("Transcribed {} samples: '{}'", audioData.size(), text);
    }

    return text;
#else
    m_logger->warn("Vosk not available, cannot transcribe");
    return "";
#endif
}

std::string VoskSTTEngine::transcribeFile(const std::string& filePath) {
#ifdef HAS_VOSK
    if (!m_available) {
        m_logger->warn("Vosk engine not available");
        return "";
    }

    m_logger->info("Transcribing file: {}", filePath);

    // Read WAV file
    std::vector<float> audioData;
    if (!readWavFile(filePath, audioData)) {
        m_logger->error("Failed to read audio file: {}", filePath);
        return "";
    }

    m_logger->debug("Read {} samples from file", audioData.size());

    // Reset recognizer for new utterance
    vosk_recognizer_reset(m_recognizer);

    // Transcribe
    std::string result = transcribe(audioData);

    // Get final result
    const char* finalResult = vosk_recognizer_final_result(m_recognizer);
    if (finalResult) {
        std::string finalText = parseVoskResult(finalResult);
        if (!finalText.empty()) {
            result = finalText;
        }
    }

    m_logger->info("Transcription result: '{}'", result);
    return result;
#else
    m_logger->warn("Vosk not available, cannot transcribe file");
    return "";
#endif
}

void VoskSTTEngine::setLanguage(const std::string& language) {
    // Vosk models are language-specific and loaded at construction time
    // Cannot change language at runtime
    m_logger->debug("Language setting ignored (Vosk model is language-specific)");
    m_logger->debug("To change language, load a different model");
}

std::string VoskSTTEngine::parseVoskResult(const std::string& jsonStr) {
    try {
        auto json = nlohmann::json::parse(jsonStr);

        // Vosk returns: {"text": "transcription"} or {"partial": "partial text"}
        if (json.contains("text")) {
            return json["text"].get<std::string>();
        } else if (json.contains("partial")) {
            return json["partial"].get<std::string>();
        }

        return "";
    } catch (const std::exception& e) {
        m_logger->error("Failed to parse Vosk result: {}", e.what());
        m_logger->debug("JSON string: {}", jsonStr);
        return "";
    }
}

bool VoskSTTEngine::readWavFile(const std::string& filePath, std::vector<float>& outSamples) {
    std::ifstream file(filePath, std::ios::binary);
    if (!file) {
        m_logger->error("Cannot open file: {}", filePath);
        return false;
    }

    // Read WAV header (44 bytes for standard PCM WAV)
    char header[44];
    file.read(header, 44);
    if (!file) {
        m_logger->error("Invalid WAV file (header too short)");
        return false;
    }

    // Verify RIFF header
    if (std::strncmp(header, "RIFF", 4) != 0) {
        m_logger->error("Not a valid WAV file (missing RIFF)");
        return false;
    }

    // Verify WAVE format
    if (std::strncmp(header + 8, "WAVE", 4) != 0) {
        m_logger->error("Not a valid WAV file (missing WAVE)");
        return false;
    }

    // Extract audio format info
    uint16_t audioFormat = *reinterpret_cast<uint16_t*>(header + 20);
    uint16_t numChannels = *reinterpret_cast<uint16_t*>(header + 22);
    uint32_t sampleRate = *reinterpret_cast<uint32_t*>(header + 24);
    uint16_t bitsPerSample = *reinterpret_cast<uint16_t*>(header + 34);
    uint32_t dataSize = *reinterpret_cast<uint32_t*>(header + 40);

    m_logger->debug("WAV format: {} channels, {} Hz, {} bits",
                    numChannels, sampleRate, bitsPerSample);

    // We expect PCM 16-bit mono
    if (audioFormat != 1) {
        m_logger->error("Unsupported audio format (expected PCM)");
        return false;
    }

    if (bitsPerSample != 16) {
        m_logger->warn("Expected 16-bit audio, got {} bits", bitsPerSample);
    }

    // Read audio data
    std::vector<int16_t> samples(dataSize / sizeof(int16_t));
    file.read(reinterpret_cast<char*>(samples.data()), dataSize);
    if (!file) {
        m_logger->error("Failed to read audio data");
        return false;
    }

    // Convert to float and handle multi-channel (take first channel)
    size_t numSamples = samples.size() / numChannels;
    outSamples.resize(numSamples);

    for (size_t i = 0; i < numSamples; ++i) {
        // Take first channel, convert int16 to float [-1.0, 1.0]
        outSamples[i] = samples[i * numChannels] / 32768.0f;
    }

    return true;
}

} // namespace celuna
