#pragma once

#include <nlohmann/json.hpp>
#include <string>
#include <regex>

using json = nlohmann::ordered_json;

//
// Bailing / Ling-Flash Function Calling Parser (arg_key / arg_value XML format)
//
// Rationale: thinking models of the Bailing family (verified: Ling-3.0-Flash)
// intermittently emit their tool call inside the thinking block instead of the
// content position:
//
//   <think>...narration...
//   <tool_call>web_search
//   <arg_key>query</arg_key>
//   <arg_value>some query</arg_value>
//   <arg_key>limit</arg_key>
//   <arg_value>3</arg_value>
//   </tool_call></think>
//
// The PEG layer consumes the thinking region as reasoning_content before the
// tool stage runs, so the call is lost: the response carries empty content,
// no structured tool calls, and the call text parked in reasoning_content.
// Reference semantic is vLLM's Ling3 parser (vllm/parser/ling3.py): <tool_call>
// acts as an implicit reasoning terminator, thinking ends where the call
// begins. This parser recovers well-formed blocks from that region. It mirrors the
// behavior validated client-side (Hermes reasoning_tool_rescue.py, 330
// rescued turns over a 2h probe with no loop breakage):
//   - only complete open-to-close blocks; truncated fragments are skipped,
//   - function name must be a bare identifier (first stripped line),
//   - argument values stay raw strings; the client coerces simple scalars
//     against its schemas downstream,
//   - never throws: malformed input yields fewer (or zero) calls.
// NOTE: a block naming the literal tag ("tool_call" as function name) is
// rescued like any other name; the client decides dispatch. Malformed
// emissions of that shape produced downstream arg errors, never loop breaks.
//
namespace bailing {

static constexpr int k_max_blocks_per_message = 8;

// Parse Bailing XML-style tool calls from free text (typically reasoning).
// Returns a JSON array of {"id": "", "type": "function",
// {"function": {"name": ..., "arguments": "<json object string>"}}.
// ids are intentionally empty: common_chat_msg::set_tool_call_ids fills them
// through the normal path when the caller appends before that call.
static json parse_tool_calls(const std::string & text) {
    json tool_calls = json::array();

    try {
        if (text.find("<tool_call>") == std::string::npos) {
            return tool_calls;
        }
        std::regex block_regex(R"(<tool_call>([\s\S]*?)</tool_call>)");
        std::regex pair_regex(
            R"(<arg_key>([\s\S]*?)</arg_key>\s*<arg_value>([\s\S]*?)</arg_value>)");
        std::regex name_regex(R"(^[A-Za-z0-9_\-]{1,64}$)");

        auto trim = [](std::string s) {
            const char * ws = " \t\n\r";
            s.erase(0, s.find_first_not_of(ws));
            if (!s.empty()) {
                s.erase(s.find_last_not_of(ws) + 1);
            }
            return s;
        };

        int blocks = 0;
        std::sregex_iterator it(text.begin(), text.end(), block_regex);
        std::sregex_iterator end;
        for (; it != end && blocks < k_max_blocks_per_message; ++it, ++blocks) {
            std::string block = (*it)[1].str();

            // Function name: first stripped line of the block.
            std::string first_line;
            {
                size_t nl = block.find('\n');
                first_line = trim(nl == std::string::npos ? block : block.substr(0, nl));
            }
            if (!std::regex_match(first_line, name_regex)) {
                continue;
            }

            json args = json::object();
            std::sregex_iterator pit(block.begin(), block.end(), pair_regex);
            for (; pit != end; ++pit) {
                std::string key = trim((*pit)[1].str());
                std::string value = trim((*pit)[2].str());
                if (!key.empty()) {
                    args[key] = value;  // last duplicate wins
                }
            }

            json tool_call = {
                {"id", ""},
                {"type", "function"},
                {"function", {
                    {"name", first_line},
                    {"arguments", args.dump()},
                }},
            };
            tool_calls.push_back(tool_call);
        }
    } catch (const std::exception &) {
        return json::array();
    }

    return tool_calls;
}

// True when the text carries at least one complete Bailing tool-call block.
static bool has_complete_block(const std::string & text) {
    try {
        std::regex block_regex(R"(<tool_call>[\s\S]*?</tool_call>)");
        return std::regex_search(text, block_regex);
    } catch (const std::exception &) {
        return false;
    }
}

// Remove executed (complete) blocks from reasoning text. Mirrors the vLLM
// Ling3 terminator semantic: reasoning ends where the call begins, so panels
// downstream never display converted calls. Only complete blocks are removed;
// truncated fragments are left untouched. Leading/trailing whitespace left by
// removal is trimmed; internal formatting is preserved.
static std::string remove_executed_blocks(const std::string & text) {
    try {
        std::regex block_regex(R"(<tool_call>[\s\S]*?</tool_call>)");
        std::string cleaned = std::regex_replace(text, block_regex, "");
        const char * ws = " \t\n\r";
        cleaned.erase(0, cleaned.find_first_not_of(ws));
        if (!cleaned.empty()) {
            cleaned.erase(cleaned.find_last_not_of(ws) + 1);
        }
        return cleaned;
    } catch (const std::exception &) {
        return text;
    }
}

}  // namespace bailing