// mm-server.cpp: HTTP server for MiniMax Music 3 generation
//
// Single binary, one port, embedded webui. The compute endpoint
// (POST /synth) is asynchronous: it validates the request, creates a
// job, pushes it to a FIFO queue, and returns the job ID immediately.
// A single worker thread processes jobs in order.
// Clients poll GET /job?id=N for status and fetch results with
// GET /job?id=N&result=1. POST /job?id=N&cancel=1 cancels a job.
//
// Job IDs are random 64-bit hex strings (non-predictable).
// Completed jobs are evicted FIFO when the pool exceeds MAX_JOBS.
// Running jobs are never evicted.
//
// Models are discovered by scanning the --models directory at startup
// (reads GGUF metadata only, no weights loaded). Module loads go
// through a ModelStore: STRICT by default (at most one coexistence
// group resident, the LM and the DiT never overlap), or resident
// forever with --keep-loaded. See the doctrine in model-store.h. GPU
// access is serialized by the single worker thread (no mutex needed).

#include "audio-io.h"
#include "model-registry.h"
#include "pipeline.h"
#include "prompt.h"
#include "request.h"
#include "version.h"
#include "yyjson.h"

// embedded webui (generated by xxd.cmake from tools/public/index.html.gz)
#include "index.html.gz.hpp"

// suppress warnings in third-party headers
#ifdef __GNUC__
#    pragma GCC diagnostic push
#    pragma GCC diagnostic ignored "-Wshadow"
#endif
#include "httplib.h"
#ifdef __GNUC__
#    pragma GCC diagnostic pop
#endif

#include <atomic>
#include <condition_variable>
#include <csignal>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <deque>
#include <functional>
#include <memory>
#include <mutex>
#include <random>
#include <string>
#include <thread>
#include <unordered_map>
#include <vector>

#ifdef _WIN32
#    include <fcntl.h>
#    include <io.h>
#    ifndef STDERR_FILENO
#        define STDERR_FILENO 2
#    endif
#else
#    include <unistd.h>
#endif

// portable fd wrappers. avoids macros that collide with C++ method names
// (e.g. sink.write() in httplib would be eaten by a write() macro).
#ifdef _WIN32
static int fd_pipe(int fd[2]) {
    return _pipe(fd, 4096, _O_BINARY);
}

static int fd_dup(int fd) {
    return _dup(fd);
}

static int fd_dup2(int src, int dst) {
    return _dup2(src, dst);
}

static int fd_read(int fd, void * buf, size_t n) {
    return _read(fd, buf, (unsigned) n);
}

static int fd_write(int fd, const void * buf, size_t n) {
    return _write(fd, buf, (unsigned) n);
}

static void fd_close(int fd) {
    _close(fd);
}
#else
static int fd_pipe(int fd[2]) {
    return pipe(fd);
}

static int fd_dup(int fd) {
    return dup(fd);
}

static int fd_dup2(int src, int dst) {
    return dup2(src, dst);
}

static int fd_read(int fd, void * buf, size_t n) {
    return (int) read(fd, buf, n);
}

static int fd_write(int fd, const void * buf, size_t n) {
    return (int) write(fd, buf, n);
}

static void fd_close(int fd) {
    close(fd);
}
#endif

// server instance pointer for the signal handler
static httplib::Server * g_svr = nullptr;

static void on_signal(int) {
    if (g_svr) {
        g_svr->stop();
    }
}

// work queue: all GPU jobs go through a single FIFO queue processed
// by one worker thread. GPU access is serialized by construction.
static std::deque<std::function<void()>> g_work_queue;
static std::mutex                        mtx_work;
static std::condition_variable           cv_work;
static bool                              g_work_stop = false;

static void work_push(std::function<void()> fn) {
    std::lock_guard<std::mutex> lock(mtx_work);
    g_work_queue.push_back(std::move(fn));
    cv_work.notify_one();
}

// worker thread: consume jobs in FIFO order until shutdown.
// on stop: finishes the current job, discards pending ones.
static void worker_main() {
    for (;;) {
        std::function<void()> fn;
        {
            std::unique_lock<std::mutex> lock(mtx_work);
            cv_work.wait(lock, [] { return g_work_stop || !g_work_queue.empty(); });
            if (g_work_stop) {
                break;
            }
            fn = std::move(g_work_queue.front());
            g_work_queue.pop_front();
        }
        fn();
    }
}

// pipeline: borrows its modules from the store stage by stage.
// only the worker thread touches it. pipeline_configure records the
// resolved paths; loads and evictions happen inside generate.
static MM3Pipeline       g_pipeline;
static MM3PipelineParams g_params;
static std::string       g_models_dir;
static bool              g_keep_loaded = false;

// model registry (populated at startup from GGUF metadata)
static ModelRegistry g_registry;

// Prompt budget of the official serving stack, enforced at submit.
// The tokenizer loads from the first LM GGUF at startup (metadata only).
static const int    MAX_PROMPT_TOKENS = 5000;
static BPETokenizer g_tok;
static bool         g_tok_ready = false;

// validate the request's model names against the registry (immutable
// after startup, safe from any thread). explicit unknown names and empty
// buckets get a 400 before the job is queued.
static bool validate_models(const MM3Request & r) {
    const struct {
        const std::vector<ModelEntry> & bucket;
        const std::string &             requested;
    } checks[] = {
        { g_registry.lm,    r.lm_model    },
        { g_registry.depth, r.depth_model },
        { g_registry.cond,  r.cond_model  },
        { g_registry.dit,   r.dit_model   },
        { g_registry.vae,   r.vae_model   },
    };

    for (const auto & c : checks) {
        if (c.bucket.empty()) {
            return false;
        }
        if (!c.requested.empty() && !registry_find(c.bucket, c.requested.c_str())) {
            return false;
        }
    }
    return true;
}

// resolve all five components against the wanted set. worker thread only
// (reads g_pipeline.wanted), called at job start so queued jobs follow
// the models requested by the previous job.
static void resolve_paths(const MM3Request & r, MM3ModelPaths & paths) {
    const MM3ModelPaths & l = g_pipeline.wanted;
    paths.lm                = registry_resolve(g_registry.lm, r.lm_model, "lm", l.lm);
    paths.depth             = registry_resolve(g_registry.depth, r.depth_model, "depth", l.depth);
    paths.cond              = registry_resolve(g_registry.cond, r.cond_model, "cond", l.cond);
    paths.dit               = registry_resolve(g_registry.dit, r.dit_model, "dit", l.dit);
    paths.vae               = registry_resolve(g_registry.vae, r.vae_model, "vae", l.vae);
}

// Build a multipart/mixed body with one JSON replay request part
// followed by its audio part, per rendered track. The boundary is
// fixed; the client splits on it and types parts by their header.
static const char * MULTIPART_BOUNDARY = "mm3-batch-boundary";

static std::string multipart_build_tracks(const std::vector<std::string> & request_parts,
                                          const std::vector<std::string> & audio_parts,
                                          const char *                     audio_mime) {
    // One set of literal fragments sizes the body exactly and builds it:
    // audio parts weigh tens of MB, growing the string through repeated
    // appends would reallocate and copy them
    const char * dash       = "--";
    const char * json_head  = "\r\nContent-Type: application/json\r\n\r\n";
    const char * audio_head = "\r\nContent-Type: ";
    const char * head_end   = "\r\n\r\n";
    const char * crlf       = "\r\n";
    const char * close_end  = "--\r\n";

    const size_t boundary_len = strlen(MULTIPART_BOUNDARY);
    const size_t per_track    = 2 * strlen(dash) + 2 * boundary_len + strlen(json_head) + 2 * strlen(crlf) +
                             strlen(audio_head) + strlen(audio_mime) + strlen(head_end);
    size_t total = strlen(dash) + boundary_len + strlen(close_end);
    for (size_t i = 0; i < audio_parts.size(); i++) {
        total += per_track + request_parts[i].size() + audio_parts[i].size();
    }

    std::string body;
    body.reserve(total);
    for (size_t i = 0; i < audio_parts.size(); i++) {
        body += dash;
        body += MULTIPART_BOUNDARY;
        body += json_head;
        body += request_parts[i];
        body += crlf;
        body += dash;
        body += MULTIPART_BOUNDARY;
        body += audio_head;
        body += audio_mime;
        body += head_end;
        body += audio_parts[i];
        body += crlf;
    }
    body += dash;
    body += MULTIPART_BOUNDARY;
    body += close_end;
    return body;
}

static const std::string MULTIPART_MIME = std::string("multipart/mixed; boundary=") + MULTIPART_BOUNDARY;

// job system: the compute endpoint creates a job and returns its ID
// immediately. the worker thread processes jobs in FIFO order, stores
// the result. the client polls GET /job?id=N until done, then fetches
// the result with GET /job?id=N&result=1.
// cancel: POST /job?id=N&cancel=1 sets the per-job flag.
enum class JobStatus : int {
    RUNNING   = 0,
    DONE      = 1,
    FAILED    = 2,
    CANCELLED = 3,
};

struct Job {
    std::string            id;
    std::atomic<JobStatus> status{ JobStatus::RUNNING };
    std::string            result_body;
    std::string            result_mime;
    std::atomic<bool>      cancel{ false };

    // memory ordering contract: result_body and result_mime are written
    // before status is stored (seq_cst). the client loads status (seq_cst)
    // and only reads result fields after seeing done/failed. this guarantees
    // visibility without an explicit mutex on the result fields.
};

static std::mutex                                            mtx_jobs;
static std::unordered_map<std::string, std::shared_ptr<Job>> g_jobs;
static std::deque<std::string>                               g_job_order;
static const int                                             MAX_JOBS = 32;

// job currently on the GPU, tracked so shutdown can cancel it and return
// within one pipeline cancel poll instead of waiting out the generation.
static std::mutex           mtx_active;
static std::shared_ptr<Job> g_active_job;

static void active_job_set(std::shared_ptr<Job> job) {
    std::lock_guard<std::mutex> lock(mtx_active);
    g_active_job = std::move(job);
}

static void active_job_cancel() {
    std::lock_guard<std::mutex> lock(mtx_active);
    if (g_active_job && g_active_job->status.load() == JobStatus::RUNNING) {
        fprintf(stderr, "[Server] Cancelling active job %s\n", g_active_job->id.c_str());
        g_active_job->cancel.store(true);
    }
}

// cancel callback for long encodes, reads the per-job atomic flag
static bool server_cancel_job(void * data) {
    auto * flag = (const std::atomic<bool> *) data;
    return flag && flag->load(std::memory_order_relaxed);
}

// generate a random hex ID (64 bits of entropy, non-predictable)
static std::string job_make_id() {
    static std::mt19937_64      rng(std::random_device{}());
    static std::mutex           mtx_rng;
    std::lock_guard<std::mutex> lock(mtx_rng);
    char                        buf[17];
    snprintf(buf, sizeof(buf), "%016llx", (unsigned long long) rng());
    return buf;
}

static std::shared_ptr<Job> job_create() {
    std::lock_guard<std::mutex> lock(mtx_jobs);
    auto                        job = std::make_shared<Job>();
    job->id                         = job_make_id();
    g_jobs[job->id]                 = job;
    g_job_order.push_back(job->id);

    // evict oldest completed jobs to stay under MAX_JOBS.
    // running jobs are never evicted.
    while ((int) g_job_order.size() > MAX_JOBS) {
        bool evicted = false;
        for (auto it = g_job_order.begin(); it != g_job_order.end(); ++it) {
            auto jit = g_jobs.find(*it);
            if (jit == g_jobs.end() || jit->second->status.load() != JobStatus::RUNNING) {
                if (jit != g_jobs.end()) {
                    g_jobs.erase(jit);
                }
                g_job_order.erase(it);
                evicted = true;
                break;
            }
        }
        if (!evicted) {
            break;
        }
    }
    return job;
}

static std::shared_ptr<Job> job_find(const std::string & id) {
    std::lock_guard<std::mutex> lock(mtx_jobs);
    auto                        it = g_jobs.find(id);
    return it != g_jobs.end() ? it->second : nullptr;
}

static const char * job_status_str(JobStatus s) {
    switch (s) {
        case JobStatus::RUNNING:
            return "running";
        case JobStatus::DONE:
            return "done";
        case JobStatus::FAILED:
            return "failed";
        case JobStatus::CANCELLED:
            return "cancelled";
    }
    return "unknown";
}

// log capture: intercept stderr via pipe, forward to terminal + ring buffer.
// SSE clients connect to /logs and receive lines in real time.
#define LOG_RING_BITS 9
#define LOG_RING_SIZE (1 << LOG_RING_BITS)
#define LOG_RING_MASK (LOG_RING_SIZE - 1)

static std::mutex              mtx_log;
static std::condition_variable cv_log;
static std::string             log_ring[LOG_RING_SIZE];
static uint64_t                log_seq = 0;

static int g_real_stderr_fd = -1;
static int g_pipe_read_fd   = -1;

static void setup_log_capture() {
    g_real_stderr_fd = fd_dup(STDERR_FILENO);
    int pipefd[2];
    if (fd_pipe(pipefd) != 0) {
        g_real_stderr_fd = -1;
        return;
    }
    g_pipe_read_fd = pipefd[0];
    fd_dup2(pipefd[1], STDERR_FILENO);
    fd_close(pipefd[1]);
}

// reader thread: drain pipe, forward to real stderr, push lines to ring.
// exits when the write end of the pipe is closed (fd_dup2 restores real stderr).
static void log_reader_main() {
    char        buf[4096];
    std::string partial;
    for (;;) {
        int n = fd_read(g_pipe_read_fd, buf, sizeof(buf));
        if (n <= 0) {
            break;
        }
        fd_write(g_real_stderr_fd, buf, (size_t) n);
        partial.append(buf, (size_t) n);
        size_t pos;
        while ((pos = partial.find('\n')) != std::string::npos) {
            std::lock_guard<std::mutex> lock(mtx_log);
            log_ring[log_seq & LOG_RING_MASK] = partial.substr(0, pos);
            log_seq++;
            cv_log.notify_all();
            partial.erase(0, pos + 1);
        }
    }
    if (!partial.empty()) {
        std::lock_guard<std::mutex> lock(mtx_log);
        log_ring[log_seq & LOG_RING_MASK] = std::move(partial);
        log_seq++;
        cv_log.notify_all();
    }
    fd_close(g_pipe_read_fd);
}

static void teardown_log_capture() {
    if (g_real_stderr_fd < 0) {
        return;
    }
    fflush(stderr);
    fd_dup2(g_real_stderr_fd, STDERR_FILENO);
    // g_real_stderr_fd stays open: the reader thread writes to it
}

// RAII: captures stderr on construction, restores + joins reader on destruction.
// safe on any exit path (early arg errors, model load failures, normal shutdown).
struct LogCapture {
    std::thread reader;

    LogCapture() {
        setup_log_capture();
        reader = std::thread(log_reader_main);
    }

    ~LogCapture() {
        teardown_log_capture();
        cv_log.notify_all();
        if (reader.joinable()) {
            reader.join();
        }

        // reader is done draining the pipe, safe to close
        if (g_real_stderr_fd >= 0) {
            fd_close(g_real_stderr_fd);
            g_real_stderr_fd = -1;
        }
    }
};

// GET /logs: SSE stream of stderr lines.
// sends backlog (up to LOG_RING_SIZE) then streams new lines in real time.
static void handle_logs(const httplib::Request &, httplib::Response & res) {
    res.set_header("Cache-Control", "no-cache");
    res.set_header("X-Accel-Buffering", "no");
    res.set_chunked_content_provider(
        "text/event-stream", [cursor = uint64_t(0), init = false](size_t, httplib::DataSink & sink) mutable -> bool {
            std::unique_lock<std::mutex> lock(mtx_log);
            if (!init) {
                uint64_t avail = log_seq < LOG_RING_SIZE ? log_seq : (uint64_t) LOG_RING_SIZE;
                cursor         = log_seq - avail;
                while (cursor < log_seq) {
                    std::string ev = "data: " + log_ring[cursor & LOG_RING_MASK] + "\n\n";
                    cursor++;
                    lock.unlock();
                    if (!sink.write(ev.c_str(), ev.size())) {
                        return false;
                    }
                    lock.lock();
                }
                init = true;
            }
            cv_log.wait_for(lock, std::chrono::seconds(2));
            while (cursor < log_seq) {
                std::string ev = "data: " + log_ring[cursor & LOG_RING_MASK] + "\n\n";
                cursor++;
                lock.unlock();
                if (!sink.write(ev.c_str(), ev.size())) {
                    return false;
                }
                lock.lock();
            }
            return true;
        });
}

// helper: set a JSON error response
static void json_error(httplib::Response & res, int code, const char * msg) {
    res.status       = code;
    std::string body = "{\"error\":\"";
    body += msg;
    body += "\"}";
    res.set_content(body, "application/json");
}

// POST /synth: validate, create job, queue the full pipeline run
static void handle_synth(const httplib::Request & req, httplib::Response & res) {
    MM3Request r;
    request_init(&r);
    if (!request_parse_json(&r, req.body.c_str())) {
        json_error(res, 400, "Malformed JSON");
        return;
    }
    if (r.caption.empty() || r.lyrics.empty()) {
        json_error(res, 400, "caption and lyrics are required");
        return;
    }
    if (r.duration <= 0.0f) {
        json_error(res, 400, "duration must be positive");
        return;
    }
    if (r.steps < 2) {
        json_error(res, 400, "steps must be at least 2");
        return;
    }
    if (r.lm_batch_size < 1 || r.lm_batch_size > g_params.max_batch) {
        json_error(res, 400, "lm_batch_size exceeds --max-batch");
        return;
    }
    if (r.synth_batch_size < 1 || r.synth_batch_size > 9) {
        json_error(res, 400, "synth_batch_size must be between 1 and 9");
        return;
    }
    if (g_tok_ready) {
        std::vector<int> ids = mm3_build_prompt_ids([](const std::string & s) { return bpe_encode(&g_tok, s, false); },
                                                    r.caption, r.lyrics);
        if ((int) ids.size() > MAX_PROMPT_TOKENS) {
            json_error(res, 400, "Prompt exceeds the 5000 token budget");
            return;
        }
    }
    if (!r.audio_codes.empty()) {
        size_t n = 1;
        for (char c : r.audio_codes) {
            if (c == ',') {
                n++;
            } else if ((c < '0' || c > '9') && c != '-') {
                json_error(res, 400, "Invalid audio_codes");
                return;
            }
        }
        if (n < 16 || n % 8 != 0) {
            json_error(res, 400, "audio_codes needs 8 codes per frame, at least 2 frames");
            return;
        }
    }

    // Output format from MM3Request.output_format. Converts the string to
    // (output_wav, wav_fmt) using the same parser the CLI uses.
    bool      output_wav = false;
    WavFormat wav_fmt    = WAV_S16;
    {
        bool is_mp3 = true;
        if (!audio_parse_format(r.output_format.c_str(), is_mp3, wav_fmt)) {
            json_error(res, 400, "Invalid output_format (use: mp3, wav16, wav24, wav32)");
            return;
        }
        output_wav = !is_mp3;
    }
    request_resolve_seed(&r);
    request_resolve_lm_seed(&r);

    if (!validate_models(r)) {
        json_error(res, 400, "Unknown model name or empty model bucket");
        return;
    }

    auto job = job_create();
    work_push([job, r, output_wav, wav_fmt]() {
        active_job_set(job);
        MM3ModelPaths paths;
        resolve_paths(r, paths);
        pipeline_configure(&g_pipeline, paths, g_params);
        fprintf(stderr, "[Server] Job %s: %s\n", job->id.c_str(), request_to_json(&r).c_str());

        std::vector<std::vector<float>> tracks;
        std::vector<std::string>        codes;
        PipelineStatus                  status = pipeline_generate(&g_pipeline, r, &job->cancel, tracks, &codes);
        if (status == PIPELINE_CANCELLED) {
            job->status.store(JobStatus::CANCELLED);
            return;
        }
        if (status != PIPELINE_OK) {
            job->status.store(JobStatus::FAILED);
            return;
        }

        // encode (peak normalize + encode), WAV_F32 preserves full range
        const char *             mime = output_wav ? "audio/wav" : "audio/mpeg";
        std::vector<std::string> parts(tracks.size());
        for (size_t i = 0; i < tracks.size(); i++) {
            std::vector<float> & audio   = tracks[i];
            int                  T_audio = (int) (audio.size() / 2);
            if (!output_wav || wav_fmt != WAV_F32) {
                audio_normalize(audio.data(), T_audio * 2, r.peak_clip);
            }
            parts[i] = output_wav ? audio_encode_wav(audio.data(), T_audio, 44100, wav_fmt) :
                                    audio_encode_mp3(audio.data(), T_audio, 44100, r.mp3_bitrate, server_cancel_job,
                                                     (void *) &job->cancel);
            if (job->cancel.load()) {
                job->status.store(JobStatus::CANCELLED);
                return;
            }
        }
        // Every response is multipart/mixed: one replay request part and
        // one audio part per track. Codes are per song, seeds per track.
        int                      M = (int) (tracks.size() / codes.size());
        std::vector<std::string> requests(parts.size());
        for (size_t i = 0; i < parts.size(); i++) {
            MM3Request replay = request_replay(r, codes[i / M], (int) (i / M), (int) (i % M));
            requests[i]       = request_to_json(&replay, true);
        }
        job->result_body = multipart_build_tracks(requests, parts, mime);
        job->result_mime = MULTIPART_MIME;
        job->status.store(JobStatus::DONE);
    });

    std::string body = "{\"id\":\"" + job->id + "\"}";
    res.set_content(body, "application/json");
}

// GET /props: version, model buckets, request defaults
static void handle_props(const httplib::Request &, httplib::Response & res) {
    yyjson_mut_doc * doc  = yyjson_mut_doc_new(NULL);
    yyjson_mut_val * root = yyjson_mut_obj(doc);
    yyjson_mut_doc_set_root(doc, root);

    yyjson_mut_obj_add_str(doc, root, "version", MM3_VERSION);

    // models: available model names per bucket
    auto add_names = [&](yyjson_mut_val * parent, const char * key, const std::vector<ModelEntry> & bucket) {
        yyjson_mut_val * arr = yyjson_mut_arr(doc);
        for (const auto & e : bucket) {
            yyjson_mut_arr_add_strncpy(doc, arr, e.name.c_str(), e.name.size());
        }
        yyjson_mut_obj_add_val(doc, parent, key, arr);
    };
    yyjson_mut_val * models = yyjson_mut_obj(doc);
    yyjson_mut_obj_add_val(doc, root, "models", models);
    add_names(models, "lm", g_registry.lm);
    add_names(models, "depth", g_registry.depth);
    add_names(models, "cond", g_registry.cond);
    add_names(models, "dit", g_registry.dit);
    add_names(models, "vae", g_registry.vae);

    // defaults: the full request schema with default values
    MM3Request def;
    request_init(&def);
    std::string  def_json = request_to_json(&def, false);
    yyjson_doc * def_doc  = yyjson_read(def_json.c_str(), def_json.size(), 0);
    if (def_doc) {
        yyjson_mut_val * copy = yyjson_val_mut_copy(doc, yyjson_doc_get_root(def_doc));
        yyjson_mut_obj_add_val(doc, root, "defaults", copy);
        yyjson_doc_free(def_doc);
    }

    char *      s   = yyjson_mut_write(doc, 0, NULL);
    std::string out = s ? s : "{}";
    free(s);
    yyjson_mut_doc_free(doc);
    res.set_content(out, "application/json");
}

static void print_usage(const char * argv0) {
    fprintf(stderr, "minimaxmusic.cpp %s\n\n", MM3_VERSION);
    fprintf(stderr,
            "Usage: %s --models <dir> [options]\n"
            "\n"
            "Required:\n"
            "  --models <dir>         Directory of GGUF model files\n"
            "\n"
            "Server:\n"
            "  --host <addr>          Listen address (default: 127.0.0.1)\n"
            "  --port <N>             Listen port (default: 8086)\n"
            "  --max-batch <N>        LM batch limit (default: 1)\n"
            "  --max-seq <N>          LM KV cache size (default: model context)\n"
            "  --keep-loaded          Keep every model resident in VRAM (default: evict between stages)\n"
            "\n"
            "Debug:\n"
            "  --no-fa                Disable flash attention\n"
            "  --no-batch-cfg         Split CFG into two separate forwards (LM + DiT)\n"
            "  --clamp-fp16           Clamp hidden states to FP16 range\n"
            "  --dump <dir>           Dump intermediate tensors\n",
            argv0);
}

int main(int argc, char ** argv) {
    std::string host = "127.0.0.1";
    int         port = 8086;

    if (argc < 2) {
        print_usage(argv[0]);
        return 1;
    }

    for (int i = 1; i < argc; i++) {
        if (strcmp(argv[i], "--host") == 0 && i + 1 < argc) {
            host = argv[++i];
        } else if (strcmp(argv[i], "--port") == 0 && i + 1 < argc) {
            port = atoi(argv[++i]);
        } else if (strcmp(argv[i], "--models") == 0 && i + 1 < argc) {
            g_models_dir = argv[++i];
        } else if (strcmp(argv[i], "--max-batch") == 0 && i + 1 < argc) {
            g_params.max_batch = atoi(argv[++i]);
            if (g_params.max_batch < 1) {
                g_params.max_batch = 1;
            }
        } else if (strcmp(argv[i], "--max-seq") == 0 && i + 1 < argc) {
            g_params.max_seq = atoi(argv[++i]);
        } else if (strcmp(argv[i], "--keep-loaded") == 0) {
            g_keep_loaded = true;
        } else if (strcmp(argv[i], "--no-fa") == 0) {
            g_params.use_fa = false;
        } else if (strcmp(argv[i], "--no-batch-cfg") == 0) {
            g_params.use_batch_cfg = false;
        } else if (strcmp(argv[i], "--clamp-fp16") == 0) {
            g_params.clamp_fp16 = true;
        } else if (strcmp(argv[i], "--dump") == 0 && i + 1 < argc) {
            g_params.dump_dir = argv[++i];
        } else {
            print_usage(argv[0]);
            return 1;
        }
    }

    if (g_models_dir.empty()) {
        fprintf(stderr, "[Server] ERROR: --models is required\n");
        print_usage(argv[0]);
        return 1;
    }

    LogCapture log_capture;

    g_pipeline.store = store_create(g_keep_loaded ? EVICT_NEVER : EVICT_STRICT);

    registry_scan(&g_registry, g_models_dir.c_str());
    if (!g_registry.lm.empty()) {
        g_tok_ready = load_bpe_from_gguf(&g_tok, g_registry.lm.front().path.c_str());
        if (!g_tok_ready) {
            fprintf(stderr, "[Server] WARNING: tokenizer load failed, prompt budget unchecked at submit\n");
        }
    }

    httplib::Server svr;
    g_svr = &svr;

    svr.set_socket_options([](socket_t sock) {
        int one = 1;
#ifdef _WIN32
        setsockopt(sock, SOL_SOCKET, SO_REUSEADDR, (const char *) &one, sizeof(one));
#else
        setsockopt(sock, SOL_SOCKET, SO_REUSEADDR, &one, sizeof(one));
#endif
    });

    // reject oversized bodies (8 MB: caption + lyrics JSON only)
    svr.set_payload_max_length(8 * 1024 * 1024);

    svr.Post("/synth", handle_synth);
    svr.Get("/health", [](const httplib::Request &, httplib::Response & res) {
        res.set_content("{\"status\":\"ok\"}", "application/json");
    });
    svr.Get("/props", handle_props);
    svr.Get("/logs", handle_logs);

    // job system endpoints
    svr.Get("/job", [](const httplib::Request & req, httplib::Response & res) {
        if (!req.has_param("id")) {
            json_error(res, 400, "Missing id parameter");
            return;
        }
        auto job = job_find(req.get_param_value("id"));
        if (!job) {
            json_error(res, 404, "Job not found");
            return;
        }
        // ?result=1: return result body
        if (req.has_param("result") && req.get_param_value("result") == "1") {
            if (job->status.load() != JobStatus::DONE) {
                json_error(res, 404, "Result not ready");
                return;
            }
            res.set_content(job->result_body, job->result_mime);
            return;
        }
        // default: return status JSON
        std::string body = "{\"status\":\"";
        body += job_status_str(job->status.load());
        body += "\"}";
        res.set_content(body, "application/json");
    });
    svr.Post("/job", [](const httplib::Request & req, httplib::Response & res) {
        if (!req.has_param("id")) {
            json_error(res, 400, "Missing id parameter");
            return;
        }
        auto job = job_find(req.get_param_value("id"));
        if (!job) {
            json_error(res, 404, "Job not found");
            return;
        }
        // ?cancel=1: cancel the job
        if (req.has_param("cancel") && req.get_param_value("cancel") == "1") {
            JobStatus status = job->status.load();
            if (status == JobStatus::RUNNING) {
                job->cancel.store(true);
                fprintf(stderr, "[Server] Cancel requested for job %s\n", job->id.c_str());
                status = JobStatus::CANCELLED;
            }
            std::string body = "{\"status\":\"";
            body += job_status_str(status);
            body += "\"}";
            res.set_content(body, "application/json");
            return;
        }
        json_error(res, 400, "Unknown action");
    });

    // embedded webui: gzipped single-page app (built by tools/webui/).
    // the browser decompresses transparently via Content-Encoding: gzip.
    // the .gz is committed to git so cloning + cmake + make gives a working UI.
    if (index_html_gz_len > 0) {
        svr.Get("/", [](const httplib::Request & req, httplib::Response & res) {
            if (req.get_header_value("Accept-Encoding").find("gzip") == std::string::npos) {
                res.set_content("Error: gzip is not supported by this browser", "text/plain");
            } else {
                res.set_header("Content-Encoding", "gzip");
                res.set_content(reinterpret_cast<const char *>(index_html_gz), index_html_gz_len,
                                "text/html; charset=utf-8");
            }
        });
    }

    // graceful shutdown on SIGINT/SIGTERM
    signal(SIGINT, on_signal);
    signal(SIGTERM, on_signal);

    // start FIFO worker thread (processes all GPU jobs in order)
    std::thread worker(worker_main);

    fprintf(stderr, "[Server] mm-server %s\n", MM3_VERSION);
    fprintf(stderr, "[Server] Listening on %s:%d\n", host.c_str(), port);
    fprintf(stderr, "[Server] Models: %zu LM, %zu Depth, %zu Cond, %zu DiT, %zu VAE\n", g_registry.lm.size(),
            g_registry.depth.size(), g_registry.cond.size(), g_registry.dit.size(), g_registry.vae.size());
    if (!svr.listen(host, port)) {
        fprintf(stderr, "[Server] FATAL: cannot bind %s:%d\n", host.c_str(), port);
    }

    // stop worker thread: cancel the active job (the pipeline polls the
    // flag between AR frames and DiT steps), discard pending ones.
    {
        std::lock_guard<std::mutex> lock(mtx_work);
        g_work_stop = true;
    }
    active_job_cancel();
    cv_work.notify_one();
    worker.join();

    store_free(g_pipeline.store);
    fprintf(stderr, "[Server] Done\n");
    return 0;
}