#include "config.h"
#include "http.h"
#include "runtime.h"
#include "parallel_runtime.h"
#include "parallel_http.h"

#include "../cli/args.h"
#include "../common/build_info.h"

#include "engine/framework/audio/resampling.h"
#include "engine/framework/core/backend.h"
#include "engine/framework/core/load_progress.h"
#include "engine/framework/debug/trace.h"

#include <csignal>
#include <filesystem>
#include <iostream>
#include <optional>
#include <stdexcept>
#include <string>
#include <utility>
#include <vector>

namespace {

volatile std::sig_atomic_t g_shutdown_requested = 0;

void request_shutdown(int) {
    g_shutdown_requested = 1;
}

bool shutdown_requested() {
    return g_shutdown_requested != 0;
}

std::optional<std::string> arg_value(int argc, char ** argv, const std::string & name) {
    for (int i = 1; i + 1 < argc; ++i) {
        if (argv[i] == name) {
            return std::string(argv[i + 1]);
        }
    }
    return std::nullopt;
}

bool has_arg(int argc, char ** argv, const std::string & name) {
    for (int i = 1; i < argc; ++i) {
        if (argv[i] == name) {
            return true;
        }
    }
    return false;
}

std::vector<std::string> arg_values(int argc, char ** argv, const std::string & name) {
    std::vector<std::string> out;
    for (int i = 1; i + 1 < argc; ++i) {
        if (argv[i] == name) {
            out.emplace_back(argv[i + 1]);
            ++i;
        }
    }
    return out;
}

std::pair<std::string, std::string> parse_key_value_arg(const std::string & text, const std::string & name) {
    const auto eq = text.find('=');
    if (eq == std::string::npos || eq == 0) {
        throw std::runtime_error(name + " must be key=value");
    }
    return {text.substr(0, eq), text.substr(eq + 1)};
}

std::filesystem::path executable_directory(const char * argv0) {
    if (argv0 == nullptr || *argv0 == '\0') {
        return std::filesystem::current_path();
    }
    std::error_code ec;
    auto path = std::filesystem::absolute(std::filesystem::path(argv0), ec);
    if (ec) {
        return std::filesystem::current_path();
    }
    path = path.lexically_normal();
    if (std::filesystem::is_regular_file(path, ec)) {
        return path.parent_path();
    }
    return std::filesystem::current_path();
}

void print_help() {
    std::cout
        << "audiocpp_server [--config <server.json>] [--ui] [--host <ip>] [--port <port>] [--backend <backend>]\n"
        << "                [--device <id>] [--list-devices [--json]] [--threads <n>] [--busy-timeout-ms <ms>]\n"
        << "                [--max-loaded-models <n>] [--idle-unload-ms <ms>] [--min-free-memory-mb <mb>]\n"
        << "                [--model-spec-override <json-or-directory|embedded>] [--voice-dir <directory>]\n"
        << "                [--frontend-listener <name>] [--frontend-option key=value]\n"
        << "                [--log] [--log-file <path>] [--load-progress 0|1]\n"
        << "                [--cors-origins <origins>] [--resample-kernel-cache-mb <mb>]\n"
        << "  --version                        print build version, commit, compiler, platform, and enabled backends\n"
        << "  --parallel-jobs                  opt into the experimental parallel runtime; per-model slots default to 1\n"
        << "  --ui                             serve the embedded WebUI\n"
        << "  --no-ui                          disable the embedded WebUI\n"
        << "  --ui-management                  allow WebUI model management and downloads; requires\n"
        << "                                   AUDIOCPP_BUILD_NATIVE_MODEL_MANAGER=ON at build time\n"
        << "  --host <ip>                      server bind address; default 127.0.0.1\n"
        << "  --port <port>                    server listening port; default 8080\n"
        << "  --backend cpu|cuda|hip|rocm|vulkan|metal|webgpu  default cuda (rocm is an alias for hip)\n"
        << "  --backend-impl <name>             optional backend implementation; default preserves existing dispatch\n"
        << "  --list-devices                   list available backend devices and exit; add --json\n"
        << "                                   for PCI ids and free/total memory as JSON\n"
        << "  --busy-timeout-ms <ms>           fail a request with 503 when the model has been\n"
        << "                                   busy this long; default 300000, 0 disables\n"
        << "  --max-loaded-models <n>          keep at most n models resident in memory, unloading\n"
        << "                                   the least recently used idle model first; 1 enforces\n"
        << "                                   a single loaded model, default 0 (no limit)\n"
        << "  --idle-unload-ms <ms>            unload all resident models after this many ms without\n"
        << "                                   any model load/run; default 0 (disabled), next request\n"
        << "                                   reloads lazily\n"
        << "  --min-free-memory-mb <mb>        refuse a model load unless host and GPU each keep at\n"
        << "                                   least this many MiB free after the load; default 0\n"
        << "                                   (guard disabled)\n"
        << "  --resample-kernel-cache-mb <mb>  keep at most this many MiB of resampling kernels cached,\n"
        << "                                   evicting the least recently used; default 32, 0 disables;\n"
        << "                                   a kernel larger than this is rebuilt on every call\n"
        << "  --voice-dir <directory>          override the shared reference voice library directory\n"
        << "  --frontend-listener <name>       opt-in frontend transport listener, e.g. https or websocket\n"
        << "  --frontend-option key=value      listener-specific option; may be repeated\n"
        << "  --cors-origins \"*\"              experimental; disabled by default. Allows browser\n"
        << "                                   requests from any origin for trusted local demos only\n"
        << "  --load-progress 0|1              with --log/--log-file, emit model-load progress keys\n"
        << "                                   (runtime.load.phase, <store>.weights.upload_progress,\n"
        << "                                   server.load.progress); default 0 (no new log lines)\n"
        << "\n"
        << "Endpoints:\n"
        << "  GET  /                           embedded WebUI (enabled by default with a config)\n"
        << "  GET  /health\n"
        << "  GET  /v1/models\n"
        << "  GET  /v1/recipes\n"
        << "  POST /v1/models/load             available with --ui-management; path, or\n"
        << "                                   package_id with the native model manager\n"
        << "  POST /v1/models/unload           available with --ui-management\n"
        << "  POST /v1/ui/upload               temporary browser upload for WebUI requests\n"
        << "  POST /v1/ui/models/install       background package download/preparation\n"
        << "  POST /v1/ui/models/install/stop  stop one active package download\n"
        << "  POST /v1/ui/models/clean-partial remove abandoned package staging files\n"
        << "  POST /v1/ui/models/delete        remove one installed package precision\n"
        << "  GET  /v1/ui/models-root          current and binary-local default models folders\n"
        << "  POST /v1/ui/models-root          select a models folder (empty path restores default)\n"
        << "  POST /v1/ui/browse-directories   list local folders for the native folder picker\n"
        << "  GET  /v1/ui/models/install-status[?id=<package>]\n"
        << "  GET  /v1/ui/models/package-sizes package sizes from metadata-only checks\n"
        << "  GET  /v1/audio/voices?model=<id>\n"
        << "  POST /v1/audio/speech\n"
        << "  POST /v1/audio/speech/live?model=<id>[&return_text=true]\n"
        << "       raw PCM in a chunked body, speech audio deltas as SSE on the same connection;\n"
        << "       return_text adds the text the model writes as text deltas\n"
        << "  POST /v1/audio/transcriptions\n"
        << "       fields: file, model, language, prompt, stream\n"
        << "  POST /v1/batches/transcriptions\n"
        << "       fields: repeated file, model, language, prompt, options\n"
        << "  POST /v1/audio/alignments\n"
        << "       fields: file, model, text, language\n"
        << "       OpenAI-style streaming: speech stream_format=sse|audio, transcription stream=true\n"
        << "  POST /v1/audio/transcriptions/live?model=<id>\n"
        << "       raw PCM in a chunked body, transcript deltas as SSE on the same connection\n"
        << "  POST /v1/tasks/run\n"
        << "  POST /v1/tasks/batch\n"
        << "  POST /v1/tasks/stream\n"
        << "       events and result in one JSON body, or as SSE with stream_format=sse\n"
        << "  POST /v1/recipes/run\n"
        << "       run a built-in multi-model recipe; set stream=true for SSE\n";
}

}  // namespace

int main(int argc, char ** argv) {
    try {
        if (has_arg(argc, argv, "--list-devices")) {
            if (has_arg(argc, argv, "--json")) {
                engine::core::print_backend_devices_json(std::cout);
            } else {
                engine::core::print_backend_devices(std::cout);
            }
            return 0;
        }
        if (has_arg(argc, argv, "--version")) {
            minitts::app::print_build_info(std::cout);
            return 0;
        }
        if (has_arg(argc, argv, "--help") || has_arg(argc, argv, "-h")) {
            print_help();
            return 0;
        }
        const bool parallel_jobs = has_arg(argc, argv, "--parallel-jobs");
        if (parallel_jobs) { engine::io::json::enable_serialized_json_parsing(); }
        const auto config_path = arg_value(argc, argv, "--config");
        const bool ui_requested = has_arg(argc, argv, "--ui");
        if (!config_path.has_value() && !ui_requested) {
            throw std::runtime_error("missing required --config argument (or use --ui for the native WebUI)");
        }
        const auto log_file = arg_value(argc, argv, "--log-file");
        engine::debug::configure_logging(engine::debug::LoggingConfig{
            has_arg(argc, argv, "--log") || log_file.has_value(),
            log_file,
        });
        // Opt-in model-load progress reporting (runtime.load.phase,
        // <store>.weights.upload_progress, server.load.progress). Default
        // off so the log for existing downstream pipelines is unchanged.
        if (const auto load_progress = arg_value(argc, argv, "--load-progress")) {
            if (*load_progress != "0" && *load_progress != "1") {
                throw std::runtime_error("--load-progress expects 0 or 1");
            }
            engine::core::set_load_progress_enabled(*load_progress == "1");
        }
        std::signal(SIGINT, request_shutdown);
        std::signal(SIGTERM, request_shutdown);
#ifdef SIGPIPE
        // Writing to a socket whose peer has already disconnected (for example a
        // client that closed an SSE/chunked stream early) would otherwise deliver
        // SIGPIPE and terminate the whole server. Ignore it so the failed send
        // surfaces as an EPIPE error on that single request thread, which
        // handle_client already unwinds cleanly, instead of taking the process down.
        std::signal(SIGPIPE, SIG_IGN);
#endif

        auto config = config_path.has_value()
            ? minitts::server::load_server_config(*config_path, parallel_jobs)
            : minitts::server::ServerConfig{};
        if (!config_path.has_value()) {
            config.lazy_load = true;
        }
        if (ui_requested) {
            config.ui_enabled = true;
            config.ui_requested = true;
        }
        if (has_arg(argc, argv, "--no-ui")) {
            config.ui_enabled = false;
            config.ui_requested = false;
        }
        if (has_arg(argc, argv, "--ui-management")) {
            config.ui_management = true;
        }
#if !defined(AUDIOCPP_HAS_NATIVE_MODEL_MANAGER)
        if (config.ui_management) {
            throw std::runtime_error(
                "UI model management is not available in this build; reconfigure with "
                "-DAUDIOCPP_BUILD_NATIVE_MODEL_MANAGER=ON");
        }
#endif
        if (const auto host = arg_value(argc, argv, "--host")) {
            config.host = *host;
        }
        if (const auto port = arg_value(argc, argv, "--port")) {
            config.port = std::stoi(*port);
        }
        if (const auto cors_origins = arg_value(argc, argv, "--cors-origins")) {
            config.cors_origins = *cors_origins;
        }
        if (const auto backend = arg_value(argc, argv, "--backend")) {
            config.backend = minitts::server::parse_server_backend(*backend);
        }
        if (const auto impl = arg_value(argc, argv, "--backend-impl")) {
            config.backend_impl = *impl;
        }
        if (const auto device = arg_value(argc, argv, "--device")) {
            config.device = std::stoi(*device);
        }
        if (const auto threads = arg_value(argc, argv, "--threads")) {
            config.threads = std::stoi(*threads);
        }
        if (const auto busy_timeout = arg_value(argc, argv, "--busy-timeout-ms")) {
            config.busy_timeout_ms = std::stoi(*busy_timeout);
        }
        if (const auto max_loaded_models = arg_value(argc, argv, "--max-loaded-models")) {
            config.max_loaded_models = std::stoi(*max_loaded_models);
        }
        if (const auto idle_unload_ms = arg_value(argc, argv, "--idle-unload-ms")) {
            config.idle_unload_ms = std::stoi(*idle_unload_ms);
        }
        if (const auto min_free_memory_mb = arg_value(argc, argv, "--min-free-memory-mb")) {
            config.min_free_memory_mb = std::stoi(*min_free_memory_mb);
        }
        if (const auto resample_kernel_cache_mb = arg_value(argc, argv, "--resample-kernel-cache-mb")) {
            config.resample_kernel_cache_mb =
                minitts::cli::parse_non_negative_int(*resample_kernel_cache_mb, "--resample-kernel-cache-mb");
        }
        if (const auto model_spec = arg_value(argc, argv, "--model-spec-override")) {
            config.model_spec_override = std::filesystem::path(*model_spec);
        }
        if (const auto voice_dir = arg_value(argc, argv, "--voice-dir")) {
            config.voice_dir = std::filesystem::path(*voice_dir);
        }
        if (const auto frontend_listener = arg_value(argc, argv, "--frontend-listener")) {
            config.frontend_listener = *frontend_listener;
        }
        for (const auto & option : arg_values(argc, argv, "--frontend-option")) {
            auto [key, value] = parse_key_value_arg(option, "--frontend-option");
            config.frontend_options[std::move(key)] = std::move(value);
        }
        if (config.frontend_listener.empty() && !config.frontend_options.empty()) {
            throw std::runtime_error("--frontend-option requires --frontend-listener");
        }
        if (!(config.cors_origins == "*" || config.cors_origins == "")) {
            throw std::runtime_error("--cors-origins must be '*' (allow all origins) or '' (disabled)");
        }
        if (config.threads <= 0) {
            throw std::runtime_error("--threads must be positive");
        }
        if (config.busy_timeout_ms < 0) {
            throw std::runtime_error("--busy-timeout-ms must be >= 0 (0 disables the guard)");
        }
        if (config.max_loaded_models < 0) {
            throw std::runtime_error("--max-loaded-models must be >= 0 (0 disables the limit)");
        }
        if (config.idle_unload_ms < 0) {
            throw std::runtime_error("--idle-unload-ms must be >= 0 (0 disables idle unload)");
        }
        if (config.min_free_memory_mb < 0) {
            throw std::runtime_error("--min-free-memory-mb must be >= 0 (0 disables the memory guard)");
        }
        if (config.resample_kernel_cache_mb.has_value()) {
            engine::audio::set_resample_kernel_cache_budget_bytes(
                static_cast<size_t>(*config.resample_kernel_cache_mb) * 1024ull * 1024ull);
        }

        const auto ui_resource_anchor = executable_directory(argc > 0 ? argv[0] : nullptr);
        auto serve = [&](auto & state) {
            if (!config.frontend_listener.empty()) {
                auto listener = state.make_frontend_listener(config.frontend_listener);
                listener->serve(
                    config.host,
                    config.port,
                    state,
                    shutdown_requested,
                    config.max_request_body_bytes,
                    config.frontend_options);
            } else {
                const auto http_listener = parallel_jobs
                    ? minitts::server::serve_parallel_http
                    : minitts::server::serve_http;
                http_listener(
                    config.host,
                    config.port,
                    state,
                    shutdown_requested,
                    config.max_request_body_bytes);
            }
        };
        if (parallel_jobs) {
            std::cerr
                << "WARNING: --parallel-jobs enables EXPERIMENTAL server behavior.\n"
                << "Validated scope: controlled CPU session/queue lifecycle tests and representative\n"
                << "CUDA/Vulkan capacity-one cold/warm inference. No catalogue multi-slot models\n"
                << "are admitted by this framework PR. Reconfiguration, unload/eviction, disconnect\n"
                << "and shutdown have partial test coverage; real concurrent GPU streams/batches,\n"
                << "Metal model execution and hung-backend cancellation remain unverified.\n"
                << "See docs/maintainers/parallel_server_behavior_contract.md for tested scopes\n"
                << "and known limitations. Busy timeout does not cancel backend execution.\n";
            minitts::server::ParallelServerState state(config, std::filesystem::current_path(), ui_resource_anchor);
            serve(state);
        } else {
            minitts::server::ServerState state(config, std::filesystem::current_path(), ui_resource_anchor);
            serve(state);
        }
        return 0;
    } catch (const std::exception & ex) {
        std::cerr << "audiocpp_server failed: " << ex.what() << "\n";
        return 1;
    }
}