/**************************************************************************************************
 * ZipVoice AXERA C++ Resident Server – main entry.
 *
 * Architecture (like Kokoro/MeloTTS resident service):
 *   1. Load models + prompts once at startup.
 *   2. Read JSON-line requests from stdin.
 *   3. For each request: build segments → sample → vocoder → write WAV → respond JSON.
 *
 * This does NOT affect the original command-line zipvoice_axera binary.
 **************************************************************************************************/

#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <iostream>
#include <sstream>
#include <string>
#include <unistd.h>

#include "zipvoice_resident_server.hpp"

// ---------------------------------------------------------------------------
// Minimal JSON parser (no external dependency).  Only enough for our wire
// protocol; NOT a general-purpose parser.
// ---------------------------------------------------------------------------

static std::string JsonGetString(const std::string& json, const std::string& key,
                                  const std::string& default_val = "") {
    std::string pattern = "\"" + key + "\":";
    auto pos = json.find(pattern);
    if (pos == std::string::npos) return default_val;
    pos += pattern.size();
    while (pos < json.size() && (json[pos] == ' ' || json[pos] == '\t')) ++pos;
    if (pos >= json.size() || json[pos] != '"') return default_val;
    ++pos;
    std::string val;
    while (pos < json.size() && json[pos] != '"') {
        if (json[pos] == '\\' && pos + 1 < json.size()) { ++pos; }
        val += json[pos]; ++pos;
    }
    return val;
}

static float JsonGetFloat(const std::string& json, const std::string& key, float default_val = 1.0f) {
    auto s = JsonGetString(json, key);
    if (s.empty()) {
        // try numeric value (not quoted)
        std::string pattern = "\"" + key + "\":";
        auto pos = json.find(pattern);
        if (pos == std::string::npos) return default_val;
        pos += pattern.size();
        while (pos < json.size() && (json[pos] == ' ' || json[pos] == '\t')) ++pos;
        return (float)std::atof(json.c_str() + pos);
    }
    return (float)std::atof(s.c_str());
}

static std::string JsonEscape(const std::string& s) {
    std::string out;
    out.reserve(s.size() + 8);
    for (char c : s) {
        switch (c) {
            case '\\': out += "\\\\"; break;
            case '"':  out += "\\\""; break;
            case '\n': out += "\\n"; break;
            case '\r': out += "\\r"; break;
            case '\t': out += "\\t"; break;
            default:   out += c; break;
        }
    }
    return out;
}

// ---------------------------------------------------------------------------
// main
// ---------------------------------------------------------------------------

int main(int argc, char** argv) {
    // Disable stdout buffering so the Python bridge sees ready/error immediately.
    setbuf(stdout, NULL);

    // Save original stdout. Redirect it to stderr during init/inference so
    // library printf() calls don't pollute the JSON protocol channel.
    int stdout_backup = dup(STDOUT_FILENO);
    dup2(STDERR_FILENO, STDOUT_FILENO);

    // Defaults (same as zipvoice.cpp command line)
    std::string root = "/data/shared/huyuan/voice_agent/001_Voice_Assistant.AXERA/models/zipvoice";
    std::string model_dir = root + "/zipvoice_distill_ax650";
    std::string token_file = root + "/resources/zipvoice_hf/zipvoice/tokens.txt";
    std::string vocoder_model = root + "/cpp/vocoder/vocos_full.axmodel";
    int max_tokens = 384, max_feat_len = 1024, num_step = 4;
    float default_speed = 1.0f;
    int seed = 42;

    // ---- parse positional args (kept simple, no getopt) ----
    for (int i = 1; i < argc; ++i) {
        std::string arg = argv[i];
        if (arg == "--root" && i + 1 < argc) root = argv[++i];
        else if (arg == "--model-dir" && i + 1 < argc) model_dir = argv[++i];
        else if (arg == "--token-file" && i + 1 < argc) token_file = argv[++i];
        else if (arg == "--vocoder-model" && i + 1 < argc) vocoder_model = argv[++i];
        else if (arg == "--max-tokens" && i + 1 < argc) max_tokens = std::atoi(argv[++i]);
        else if (arg == "--max-feat-len" && i + 1 < argc) max_feat_len = std::atoi(argv[++i]);
        else if (arg == "--num-step" && i + 1 < argc) num_step = std::atoi(argv[++i]);
        else if (arg == "--speed" && i + 1 < argc) default_speed = (float)std::atof(argv[++i]);
        else if (arg == "--seed" && i + 1 < argc) seed = std::atoi(argv[++i]);
    }

    ZipVoiceResidentServer server;
    bool ok = server.Init(root, model_dir, token_file, vocoder_model,
                          max_tokens, max_feat_len, num_step, default_speed, seed);

    // Restore stdout for JSON protocol
    dup2(stdout_backup, STDOUT_FILENO);
    close(stdout_backup);
    setbuf(stdout, NULL);

    if (!ok) {
        std::cout << "{\"ready\":false,\"error\":\"init failed\"}" << std::endl;
        return 1;
    }

    std::cout << "{\"ready\":true}" << std::endl;
    std::cout.flush();

    std::string line;
    while (std::getline(std::cin, line)) {
        if (line.empty()) continue;
        if (line == "quit" || line == "exit") break;

        // Handle control commands
        std::string cmd_type = JsonGetString(line, "cmd");
        if (cmd_type == "cache_prompt") {
            std::string key = JsonGetString(line, "key", "custom");
            std::string wav = JsonGetString(line, "wav");
            std::string txt = JsonGetString(line, "text");
            bool ok = server.CachePrompt(key, wav, txt);
            std::cout << "{\"id\":\"" << JsonGetString(line, "id", "0")
                      << "\",\"cmd\":\"cache_prompt\",\"ok\":" << (ok ? "true" : "false")
                      << ",\"key\":\"" << key << "\"}" << std::endl;
            continue;
        }
        if (cmd_type == "prompts") {
            std::cout << "{\"id\":\"" << JsonGetString(line, "id", "0")
                      << "\",\"cmd\":\"prompts\",\"ok\":true,\"prompts\":[\"zh\",\"en\",\"custom\"]}" << std::endl;
            std::cout.flush();
            continue;
        }

        // Parse request: {"id":"1","text":"你好","prompt":"zh","speed":1.0,"output_wav":"/tmp/x.wav"}
        ZipVoiceRequest req;
        req.id = JsonGetString(line, "id", "0");
        req.text = JsonGetString(line, "text");
        req.prompt = JsonGetString(line, "prompt", "zh");
        req.speed = JsonGetFloat(line, "speed", default_speed);
        req.output_wav = JsonGetString(line, "output_wav", "/tmp/zipvoice_daemon.wav");

        if (req.text.empty()) {
            std::cout << "{\"id\":\"" << req.id << "\",\"ok\":false,\"error\":\"empty text\"}" << std::endl;
            std::cout.flush();
            continue;
        }

        // Redirect library printf() to stderr during inference so they
        // don't pollute the JSON stdout channel.
        int saved_stdout = dup(STDOUT_FILENO);
        dup2(STDERR_FILENO, STDOUT_FILENO);
        auto resp = server.Infer(req);
        dup2(saved_stdout, STDOUT_FILENO);
        close(saved_stdout);

        // Emit JSON response
        char buf[1024];
        snprintf(buf, sizeof(buf),
                 "{\"id\":\"%s\",\"ok\":%s,\"path\":\"%s\",\"wall_sec\":%.3f,\"model_sec\":%.3f,\"rtf\":%.3f%s%s}",
                 resp.id.c_str(),
                 resp.ok ? "true" : "false",
                 JsonEscape(resp.ok ? resp.path : "").c_str(),
                 resp.wall_sec, resp.model_sec, resp.rtf,
                 resp.ok ? "" : ",\"error\":\"",
                 resp.ok ? "" : JsonEscape(resp.error).c_str());
        std::cout << buf;
        if (!resp.ok) std::cout << "\"";
        std::cout << std::endl;
        std::cout.flush();
    }

    return 0;
}