#include "lemon/backends/whisper_server.h"
#include "lemon/backends/backend_utils.h"
#include "lemon/backend_manager.h"
#include "lemon/runtime_config.h"
#include "lemon/audio_types.h"
#include "lemon/utils/custom_args.h"
#include "lemon/utils/http_client.h"
#include "lemon/utils/process_manager.h"
#include "lemon/error_types.h"
#include <iostream>
#include <filesystem>
#include <fstream>
#include <random>
#include <sstream>
#include <iomanip>
#include <set>
#include <vector>
#include <lemon/utils/aixlog.hpp>

#ifdef _WIN32
#include <windows.h>
#else
#include <sys/stat.h>
#include <unistd.h>
#endif

namespace fs = std::filesystem;
using namespace lemon::utils;

namespace lemon {
namespace backends {

WhisperServer::WhisperServer(const std::string& log_level, ModelManager* model_manager, BackendManager* backend_manager)
    : WrappedServer("whisper-server", log_level, model_manager, backend_manager) {

    // Create temp directory for audio files
    temp_dir_ = fs::temp_directory_path() / "lemonade_audio";
    fs::create_directories(temp_dir_);
}

WhisperServer::~WhisperServer() {
    unload();

    // Clean up temp directory
    try {
        if (fs::exists(temp_dir_)) {
            fs::remove_all(temp_dir_);
        }
    } catch (const std::exception& e) {
        LOG(WARNING, "WhisperServer") << "Could not clean up temp directory: "
                  << e.what() << std::endl;
    }
}

InstallParams WhisperServer::get_install_params(const std::string& backend, const std::string& version) {
    InstallParams params;

    if (backend == "npu") {
        params.repo = "lemonade-sdk/whisper.cpp-builds";
#ifdef _WIN32
        params.filename = "whisper-" + version + "-windows-npu-x64.zip";
#else
        throw std::runtime_error("NPU whisper.cpp only supported on Windows");
#endif
    } else if (backend == "cpu") {
#ifdef _WIN32
        params.repo = "ggml-org/whisper.cpp";
        params.filename = "whisper-bin-x64.zip";
#elif defined(__linux__)
        params.repo = "lemonade-sdk/whisper.cpp-builds";
        params.filename = "whisper-" + version + "-linux-cpu-x86_64.tar.gz";
#elif defined(__APPLE__)
        params.repo = "ggml-org/whisper.cpp";
        params.filename = "whisper-bin-arm64.zip";
#else
        throw std::runtime_error("Unsupported platform for whisper.cpp");
#endif
    } else if (backend == "vulkan") {
#if defined(__linux__)
        params.repo = "lemonade-sdk/whisper.cpp-builds";
        params.filename = "whisper-" + version + "-linux-vulkan-x86_64.tar.gz";
#else
        throw std::runtime_error("Vulkan whisper.cpp backend is currently supported only on Linux");
#endif
    } else {
        throw std::runtime_error("[WhisperServer] Unknown whisper backend: " + backend);
    }

    return params;
}

// Helper to determine NPU compiled cache info based on model info from server_models.json
static std::pair<std::string, std::string> get_npu_cache_info(const ModelInfo& model_info) {
    std::string npu_cache = model_info.checkpoint("npu_cache");
    std::string npu_cache_repo = "";
    std::string npu_cache_filename = "";

    if (!npu_cache.empty()) {
        size_t colon_pos = npu_cache.find(':');
        if (colon_pos != std::string::npos) {
            npu_cache_repo = npu_cache.substr(0, colon_pos);
            npu_cache_filename = npu_cache.substr(colon_pos + 1);
        }
    }

    if (!npu_cache_repo.empty() && !npu_cache_filename.empty()) {
        LOG(INFO, "WhisperServer") << "Using NPU cache from server_models.json: "
                  << npu_cache_repo << " / " << npu_cache_filename << std::endl;
        return {npu_cache_repo, npu_cache_filename};
    }

    // No NPU cache configured for this model in server_models.json
    LOG(INFO, "WhisperServer") << "No NPU cache configured for model: " << model_info.model_name << std::endl;
    return {"", ""};
}

// Helper to download NPU compiled cache (.rai file) for a given ggml .bin model
void WhisperServer::download_npu_compiled_cache(const std::string& model_path,
                                                const ModelInfo& model_info,
                                                bool do_not_upgrade) {
    auto [cache_repo, cache_filename] = get_npu_cache_info(model_info);

    if (cache_repo.empty() || cache_filename.empty()) {
        LOG(INFO, "WhisperServer") << "No NPU compiled cache available for this model" << std::endl;
        return;
    }

    LOG(INFO, "WhisperServer") << "Downloading NPU compiled cache: " << cache_filename << std::endl;
    LOG(INFO, "WhisperServer") << "From repository: " << cache_repo << std::endl;

    // Determine where to place the .rai file (must be in the same directory as .bin file)
    fs::path model_dir = fs::path(model_path).parent_path();
    fs::path cache_path = model_dir / cache_filename;

    // Check if cache already exists
    if (fs::exists(cache_path) && !do_not_upgrade) {
        LOG(INFO, "WhisperServer") << "NPU cache already exists: " << cache_path << std::endl;
        return;
    }

    try {
        // Download .rai file directly from HuggingFace using HttpClient
        std::string hf_url = "https://huggingface.co/" + cache_repo + "/resolve/main/" + cache_filename;

        LOG(INFO, "WhisperServer") << "Downloading from: " << hf_url << std::endl;

        // Download directly to the target location
        auto download_result = utils::HttpClient::download_file(
            hf_url,
            cache_path.string(),
            utils::create_throttled_progress_callback()
        );

        if (!download_result.success) {
            throw std::runtime_error("Failed to download NPU cache from: " + hf_url + " - " + download_result.error_message);
        }

        LOG(INFO, "WhisperServer") << "NPU cache ready at: " << cache_path << std::endl;

    } catch (const std::exception& e) {
        if (fs::exists(cache_path)) {
            fs::remove(cache_path);
            LOG(INFO, "WhisperServer") << "Cleaned up partial cache file" << std::endl;
        }

        LOG(WARNING, "WhisperServer") << "Failed to download NPU cache: " << e.what() << std::endl;
        LOG(WARNING, "WhisperServer") << "Continuing without NPU cache (may cause runtime errors)" << std::endl;
    }
}

void WhisperServer::load(const std::string& model_name,
                        const ModelInfo& model_info,
                        const RecipeOptions& options,
                        bool do_not_upgrade) {
    LOG(INFO, "WhisperServer") << "Loading model: " << model_name << std::endl;
    LOG(INFO, "WhisperServer") << "Per-model settings: " << options.to_log_string() << std::endl;

    std::string whispercpp_backend = options.get_option("whispercpp_backend");
    std::string whispercpp_args = options.get_option("whispercpp_args");

    RuntimeConfig::validate_backend_choice("whispercpp", whispercpp_backend);

    // Update device type based on the actual backend selected.
    // get_device_type_from_recipe() defaults whispercpp to CPU, but npu/vulkan use different devices.
    if (whispercpp_backend == "npu") {
        device_type_ = DEVICE_NPU;
    } else if (whispercpp_backend == "vulkan") {
        device_type_ = DEVICE_GPU;
    } else {
        device_type_ = DEVICE_CPU;
    }

    backend_manager_->install_backend(SPEC.recipe, whispercpp_backend);

    std::string model_path = model_info.resolved_path();
    if (model_path.empty()) {
        throw std::runtime_error("Model file not found for checkpoint: " + model_info.checkpoint());
    }

    LOG(INFO, "WhisperServer") << "Using model: " << model_path << std::endl;
    LOG(INFO, "WhisperServer") << "Using backend: " << whispercpp_backend << std::endl;

    // For NPU backend, download the compiled cache (.rai file). This is a must-have for NPU backend.
    if (whispercpp_backend == "npu") {
        download_npu_compiled_cache(model_path, model_info, do_not_upgrade);
    }

    // Get whisper-server executable path
    std::string exe_path = BackendUtils::get_backend_binary_path(SPEC, whispercpp_backend);

    // Choose a port
    port_ = choose_port();
    if (port_ == 0) {
        throw std::runtime_error("Failed to find an available port");
    }

    LOG(INFO, "WhisperServer") << "Starting server on port " << port_ << std::endl;

    // Build command line arguments. Lemonade manages the model path and port;
    // optional whisper-server flags like --convert come from whispercpp_args.
    // Note: Don't include exe_path here - ProcessManager::start_process already handles it
    std::vector<std::string> args = {
        "-m", model_path,
        "--port", std::to_string(port_)
    };

    std::set<std::string> reserved_flags = {
        "-m",
        "--model",
        "--port"
    };

    if (!whispercpp_args.empty()) {
        std::string validation_error = validate_custom_args(whispercpp_args, reserved_flags);
        if (!validation_error.empty()) {
            throw std::invalid_argument(
                "Invalid custom whisper-server arguments:\n" + validation_error
            );
        }

        LOG(DEBUG, "WhisperServer") << "Adding custom arguments: " << whispercpp_args << std::endl;
        std::vector<std::string> custom_args_vec = parse_custom_args(whispercpp_args);
        args.insert(args.end(), custom_args_vec.begin(), custom_args_vec.end());
    }

    // Note: whisper-server doesn't support --debug flag

    // Set up environment variables for shared library loading
    std::vector<std::pair<std::string, std::string>> env_vars;
    fs::path exe_dir = fs::path(exe_path).parent_path();

#ifndef _WIN32
    // set LD_LIBRARY_PATH to include executable directory
    std::string lib_path = exe_dir.string();

    const char* existing_ld_path = std::getenv("LD_LIBRARY_PATH");
    if (existing_ld_path && strlen(existing_ld_path) > 0) {
        lib_path = lib_path + ":" + std::string(existing_ld_path);
    }

    env_vars.push_back({"LD_LIBRARY_PATH", lib_path});
    if (is_debug()) {
        std::cout << "[WhisperServer] Setting LD_LIBRARY_PATH=" << lib_path << std::endl;
    }
#endif

    // Launch the subprocess
    process_handle_ = utils::ProcessManager::start_process(
        exe_path,
        args,
        "",     // working_dir (empty = current)
        is_debug(),  // inherit_output
        false,  // filter_health_logs
        env_vars
    );

    if (process_handle_.pid == 0) {
        throw std::runtime_error("Failed to start whisper-server process");
    }

    LOG(INFO, "WhisperServer") << "Process started with PID: " << process_handle_.pid << std::endl;

    // Wait for server to be ready
    if (!wait_for_ready("/health")) {
        unload();
        throw std::runtime_error("whisper-server failed to start or become ready");
    }

    LOG(INFO, "WhisperServer") << "Server is ready!" << std::endl;
}

void WhisperServer::unload() {
    if (process_handle_.pid != 0) {
        LOG(INFO, "WhisperServer") << "Stopping server (PID: " << process_handle_.pid << ")" << std::endl;
        utils::ProcessManager::stop_process(process_handle_);
        process_handle_ = {nullptr, 0};
        port_ = 0;
    }
}

// ICompletionServer implementation - not supported for Whisper
json WhisperServer::chat_completion(const json& request) {
    return json{
        {"error", {
            {"message", "Whisper models do not support chat completion. Use audio transcription endpoints instead."},
            {"type", "unsupported_operation"},
            {"code", "model_not_applicable"}
        }}
    };
}

json WhisperServer::completion(const json& request) {
    return json{
        {"error", {
            {"message", "Whisper models do not support text completion. Use audio transcription endpoints instead."},
            {"type", "unsupported_operation"},
            {"code", "model_not_applicable"}
        }}
    };
}

json WhisperServer::responses(const json& request) {
    return json{
        {"error", {
            {"message", "Whisper models do not support responses. Use audio transcription endpoints instead."},
            {"type", "unsupported_operation"},
            {"code", "model_not_applicable"}
        }}
    };
}

// Audio file handling helpers
std::string WhisperServer::save_audio_to_temp(const std::string& audio_data,
                                              const std::string& filename) {
    // Generate unique filename
    std::random_device rd;
    std::mt19937 gen(rd());
    std::uniform_int_distribution<> dis(0, 999999);

    std::string ext = fs::path(filename).extension().string();
    if (ext.empty()) {
        ext = ".audio";  // Default extension
    }

    std::stringstream ss;
    ss << "audio_" << std::setfill('0') << std::setw(6) << dis(gen) << ext;

    fs::path temp_file = temp_dir_ / ss.str();

    // Write audio data to file
    std::ofstream outfile(temp_file, std::ios::binary);
    if (!outfile) {
        throw std::runtime_error("Failed to create temporary audio file: " + temp_file.string());
    }

    outfile.write(audio_data.data(), audio_data.size());
    outfile.close();

    LOG(DEBUG, "WhisperServer") << "Saved audio to temp file: " << temp_file << std::endl;

    return temp_file.string();
}

void WhisperServer::cleanup_temp_file(const std::string& path) {
    try {
        if (fs::exists(path)) {
            fs::remove(path);
            LOG(DEBUG, "WhisperServer") << "Cleaned up temp file: " << path << std::endl;
        }
    } catch (const std::exception& e) {
        LOG(WARNING, "WhisperServer") << "Could not delete temp file " << path
                  << ": " << e.what() << std::endl;
    }
}

void WhisperServer::validate_audio_file(const std::string& path) {
    if (!fs::exists(path)) {
        throw std::runtime_error("Audio file does not exist: " + path);
    }

    std::uintmax_t file_size = fs::file_size(path);
    if (file_size > audio::Limits::MAX_FILE_SIZE_BYTES) {
        throw std::runtime_error("Audio file exceeds maximum size of 25MB");
    }

    if (file_size == 0) {
        throw std::runtime_error("Audio file is empty");
    }
}

json WhisperServer::build_transcription_request(const json& request, bool translate) {
    json whisper_req;

    // Required fields
    if (request.contains("file_path")) {
        whisper_req["file"] = request["file_path"];
    }

    // Optional fields
    if (request.contains("language") && !translate) {
        // For transcription, respect language hint
        whisper_req["language"] = request["language"];
    }
    // For translation, don't specify language (always translates to English)

    if (request.contains("prompt")) {
        whisper_req["prompt"] = request["prompt"];
    }

    if (request.contains("temperature")) {
        whisper_req["temperature"] = request["temperature"];
    }

    if (request.contains("response_format")) {
        whisper_req["response_format"] = request["response_format"];
    } else {
        whisper_req["response_format"] = "json";  // Default
    }

    // Add translate flag if needed
    if (translate) {
        whisper_req["translate"] = true;
    }

    return whisper_req;
}

// Forward audio file to whisper-server using multipart form-data
json WhisperServer::forward_multipart_audio_request(const std::string& file_path,
                                                    const json& params,
                                                    bool translate) {
    // Read the audio file content
    std::ifstream file(file_path, std::ios::binary);
    if (!file) {
        throw std::runtime_error("Could not open audio file: " + file_path);
    }

    std::ostringstream oss;
    oss << file.rdbuf();
    std::string file_content = oss.str();
    file.close();

    LOG(DEBUG, "WhisperServer") << "Audio file size: " << file_content.size() << " bytes" << std::endl;

    // Determine content type based on file extension
    fs::path filepath(file_path);
    std::string ext = filepath.extension().string();
    std::string content_type = "audio/wav";  // Default

    // Map common audio extensions to MIME types
    if (ext == ".mp3") content_type = "audio/mpeg";
    else if (ext == ".wav") content_type = "audio/wav";
    else if (ext == ".m4a") content_type = "audio/mp4";
    else if (ext == ".ogg") content_type = "audio/ogg";
    else if (ext == ".flac") content_type = "audio/flac";
    else if (ext == ".webm") content_type = "audio/webm";

    std::vector<utils::MultipartField> fields;

    utils::MultipartField audio_file;
    audio_file.name = "file";
    audio_file.data = file_content;
    audio_file.filename = filepath.filename().string();
    audio_file.content_type = content_type;
    fields.push_back(audio_file);

    // Add optional parameters as form fields
    std::string response_format = params.value("response_format", "json");
    utils::MultipartField fmt_field;
    fmt_field.name = "response_format";
    fmt_field.data = response_format;
    fields.push_back(fmt_field);

    utils::MultipartField temp_field;
    temp_field.name = "temperature";
    if (params.contains("temperature")) {
        temp_field.data = std::to_string(params["temperature"].get<double>());
    } else {
        temp_field.data = "0.0";
    }
    fields.push_back(temp_field);

    if (params.contains("language")) {
        utils::MultipartField lang_field;
        lang_field.name = "language";
        lang_field.data = params["language"].get<std::string>();
        fields.push_back(lang_field);
    }

    if (params.contains("prompt")) {
        utils::MultipartField prompt_field;
        prompt_field.name = "prompt";
        prompt_field.data = params["prompt"].get<std::string>();
        fields.push_back(prompt_field);
    }

    if (translate) {
        utils::MultipartField translate_field;
        translate_field.name = "translate";
        translate_field.data = "true";
        fields.push_back(translate_field);
    }

    const std::string url = "http://127.0.0.1:" + std::to_string(port_) + "/inference";
    LOG(DEBUG, "WhisperServer") << "Sending multipart request to " << url << std::endl;

    auto res = utils::HttpClient::post_multipart(url, fields, 300);

    LOG(DEBUG, "WhisperServer") << "Response status: " << res.status_code << std::endl;
    LOG(DEBUG, "WhisperServer") << "Response body: " << res.body << std::endl;

    if (res.status_code != 200) {
        throw std::runtime_error("whisper-server returned status " +
                                std::to_string(res.status_code) + ": " + res.body);
    }

    // Try to parse as JSON
    try {
        return json::parse(res.body);
    } catch (const json::parse_error&) {
        // If response_format is not json, return it wrapped
        return json{{"text", res.body}};
    }
}

json WhisperServer::forward_multipart_audio_data(const std::string& audio_data,
                                                  const std::string& filename,
                                                  const json& params,
                                                  bool translate) {
    if (audio_data.empty()) {
        throw std::runtime_error("Empty audio data");
    }

    LOG(DEBUG, "WhisperServer") << "Audio data size: " << audio_data.size() << " bytes (no file I/O)" << std::endl;

    // Determine content type based on filename extension
    fs::path filepath(filename);
    std::string ext = filepath.extension().string();
    std::string content_type = "audio/wav";  // Default

    if (ext == ".mp3") content_type = "audio/mpeg";
    else if (ext == ".wav") content_type = "audio/wav";
    else if (ext == ".m4a") content_type = "audio/mp4";
    else if (ext == ".ogg") content_type = "audio/ogg";
    else if (ext == ".flac") content_type = "audio/flac";
    else if (ext == ".webm") content_type = "audio/webm";

    std::vector<utils::MultipartField> fields;

    utils::MultipartField audio_file;
    audio_file.name = "file";
    audio_file.data = audio_data;
    audio_file.filename = filepath.filename().string();
    audio_file.content_type = content_type;
    fields.push_back(audio_file);

    std::string response_format = params.value("response_format", "json");
    utils::MultipartField fmt_field;
    fmt_field.name = "response_format";
    fmt_field.data = response_format;
    fields.push_back(fmt_field);

    utils::MultipartField temp_field;
    temp_field.name = "temperature";
    temp_field.data = params.contains("temperature")
        ? std::to_string(params["temperature"].get<double>())
        : "0.0";
    fields.push_back(temp_field);

    if (params.contains("language")) {
        utils::MultipartField lang_field;
        lang_field.name = "language";
        lang_field.data = params["language"].get<std::string>();
        fields.push_back(lang_field);
    }

    if (params.contains("prompt")) {
        utils::MultipartField prompt_field;
        prompt_field.name = "prompt";
        prompt_field.data = params["prompt"].get<std::string>();
        fields.push_back(prompt_field);
    }

    if (translate) {
        utils::MultipartField translate_field;
        translate_field.name = "translate";
        translate_field.data = "true";
        fields.push_back(translate_field);
    }

    const std::string url = "http://127.0.0.1:" + std::to_string(port_) + "/inference";
    LOG(DEBUG, "WhisperServer") << "Sending multipart request to " << url << " (direct data)" << std::endl;

    auto res = utils::HttpClient::post_multipart(url, fields, 300);

    LOG(DEBUG, "WhisperServer") << "Response status: " << res.status_code << std::endl;

    if (res.status_code != 200) {
        throw std::runtime_error("whisper-server returned status " +
                                std::to_string(res.status_code) + ": " + res.body);
    }

    try {
        return json::parse(res.body);
    } catch (const json::parse_error&) {
        return json{{"text", res.body}};
    }
}

// IAudioServer implementation
json WhisperServer::audio_transcriptions(const json& request) {
    try {
        // Extract audio data from request
        if (!request.contains("file_data")) {
            throw std::runtime_error("Missing 'file_data' in request");
        }

        std::string audio_data = request["file_data"].get<std::string>();
        std::string filename = request.value("filename", "audio.wav");

        // Send directly to whisper-server without file I/O
        return forward_multipart_audio_data(audio_data, filename, request, false);

    } catch (const std::exception& e) {
        return json{
            {"error", {
                {"message", std::string("Transcription failed: ") + e.what()},
                {"type", "audio_processing_error"}
            }}
        };
    }
}

} // namespace backends
} // namespace lemon
