-- Quick math — offline-first arithmetic evaluator.
--
-- Strategy:
--   1. Strip the trigger phrase from the user's utterance.
--   2. Normalise Russian/English operator words (плюс/minus/в степени/...) into
--      symbols so the remainder looks like a regular numeric expression.
--   3. Try a small shunting-yard parser on the result. If it succeeds and the
--      output is finite, speak it INSTANTLY — no LLM round-trip.
--   4. Only fall back to the LLM if the local parse failed or returned NaN.
--      That covers word problems, multi-step equations, unit conversions.
--
-- This makes 95% of "сколько будет X плюс Y" complete in under 50ms regardless
-- of network or GROQ_TOKEN.

local lang = jarvis.context.language
local raw = (jarvis.context.phrase or ""):lower()

local triggers = {
    "сколько будет", "сколько это", "посчитай", "вычисли", "реши",
    "how much is", "calculate", "compute", "solve", "what is",
}
local query = jarvis.text.strip_trigger(raw, triggers)
if not query or query == "" then query = raw end
query = query:gsub("^%s+", ""):gsub("%s+$", "")

if query == "" then
    return jarvis.cmd.not_found(lang == "ru" and "Что посчитать?" or "What to compute?")
end

-- ── 1. Russian/English operator → symbol normalisation ────────────────────

local function normalise(s)
    local mapping = {
        { "плюс",                                "+" },
        { "прибавить к?",                        "+" },
        { "минус",                               "-" },
        { "отнять",                              "-" },
        { "вычесть",                             "-" },
        { "умножить на",                         "*" },
        { "умножь на",                           "*" },
        { "помножить на",                        "*" },
        { "разделить на",                        "/" },
        { "поделить на",                         "/" },
        { "делить на",                           "/" },
        { "в степени",                           "^" },
        { "возвести в степень",                  "^" },
        { "остаток от деления на",               "%%" },
        { "по модулю",                           "%%" },
        { "квадрат",                             "^2" },
        { "куб",                                 "^3" },
        { "plus",                                "+" },
        { "minus",                               "-" },
        { "times",                               "*" },
        { "multiplied by",                       "*" },
        { "divided by",                          "/" },
        { "to the power of",                     "^" },
        { "squared",                             "^2" },
        { "cubed",                               "^3" },
        { "modulo",                              "%%" },
        { "mod",                                 "%%" },
    }
    for _, pair in ipairs(mapping) do
        s = s:gsub(pair[1], pair[2])
    end
    -- comma as decimal separator (русская запись)
    s = s:gsub("(%d),(%d)", "%1.%2")
    return s
end

-- Pattern: "корень из X" / "квадратный корень из X" → math.sqrt
local function handle_sqrt(s)
    local n = s:match("^%s*к?орен?ь?%s*и?з?%s+(%-?[%d%.]+)%s*$")
    if not n then n = s:match("^%s*sqrt%s*%(?%s*(%-?[%d%.]+)%s*%)?%s*$") end
    if n then
        local v = tonumber(n)
        if v and v >= 0 then return math.sqrt(v) end
    end
    return nil
end

-- Pattern: "X процентов от Y" / "X percent of Y" → X * Y / 100
local function handle_percent(s)
    local x, y = s:match("^%s*(%-?[%d%.]+)%s*процент[%a]*%s*от%s+(%-?[%d%.]+)%s*$")
    if not x then
        x, y = s:match("^%s*(%-?[%d%.]+)%s*percent%s*of%s+(%-?[%d%.]+)%s*$")
    end
    if x and y then
        local fx, fy = tonumber(x), tonumber(y)
        if fx and fy then return fx * fy / 100 end
    end
    return nil
end

-- ── 2. Tiny shunting-yard for "<num> <op> <num> <op> <num> ..." with parens
local function tokenize(s)
    local tokens = {}
    local i = 1
    while i <= #s do
        local c = s:sub(i, i)
        if c:match("[%s]") then
            i = i + 1
        elseif c:match("[%d%.]") then
            local j = i
            while j <= #s and s:sub(j, j):match("[%d%.]") do j = j + 1 end
            table.insert(tokens, { kind = "num", val = tonumber(s:sub(i, j - 1)) })
            i = j
        elseif c == "+" or c == "-" or c == "*" or c == "/" or c == "^" or c == "%" then
            -- handle unary minus / plus at expression start or after operator/paren
            if (c == "-" or c == "+") and (#tokens == 0 or tokens[#tokens].kind == "op" or tokens[#tokens].kind == "lparen") then
                local j = i + 1
                while j <= #s and s:sub(j, j):match("[%s]") do j = j + 1 end
                if j <= #s and s:sub(j, j):match("[%d%.]") then
                    local k = j
                    while k <= #s and s:sub(k, k):match("[%d%.]") do k = k + 1 end
                    local v = tonumber(s:sub(j, k - 1))
                    if v then
                        table.insert(tokens, { kind = "num", val = (c == "-" and -v or v) })
                        i = k
                    else
                        return nil
                    end
                else
                    return nil
                end
            else
                table.insert(tokens, { kind = "op", val = c })
                i = i + 1
            end
        elseif c == "(" then
            table.insert(tokens, { kind = "lparen" }); i = i + 1
        elseif c == ")" then
            table.insert(tokens, { kind = "rparen" }); i = i + 1
        else
            return nil -- unknown char
        end
    end
    return tokens
end

local function precedence(op)
    if op == "+" or op == "-" then return 1 end
    if op == "*" or op == "/" or op == "%" then return 2 end
    if op == "^" then return 3 end
    return 0
end

local function eval_op(a, b, op)
    if op == "+" then return a + b end
    if op == "-" then return a - b end
    if op == "*" then return a * b end
    if op == "/" then if b == 0 then return nil end; return a / b end
    if op == "%" then if b == 0 then return nil end; return a % b end
    if op == "^" then return a ^ b end
    return nil
end

local function shunting_yard_eval(tokens)
    if not tokens or #tokens == 0 then return nil end
    local values, ops = {}, {}

    local function apply()
        local op = table.remove(ops)
        local b = table.remove(values)
        local a = table.remove(values)
        if a == nil or b == nil then return false end
        local r = eval_op(a, b, op)
        if r == nil then return false end
        table.insert(values, r)
        return true
    end

    for _, t in ipairs(tokens) do
        if t.kind == "num" then
            table.insert(values, t.val)
        elseif t.kind == "op" then
            while #ops > 0 and ops[#ops] ~= "(" and precedence(ops[#ops]) >= precedence(t.val) do
                if not apply() then return nil end
            end
            table.insert(ops, t.val)
        elseif t.kind == "lparen" then
            table.insert(ops, "(")
        elseif t.kind == "rparen" then
            while #ops > 0 and ops[#ops] ~= "(" do
                if not apply() then return nil end
            end
            if ops[#ops] ~= "(" then return nil end
            table.remove(ops)
        end
    end
    while #ops > 0 do
        if ops[#ops] == "(" then return nil end
        if not apply() then return nil end
    end
    if #values ~= 1 then return nil end
    return values[1]
end

-- ── 3. Format result for speech (no scientific notation, trim trailing zeros)
local function format_result(n)
    if n ~= n then return nil end                       -- NaN
    if n == math.huge or n == -math.huge then return nil end
    if math.abs(n - math.floor(n)) < 1e-9 then
        return tostring(math.floor(n + (n < 0 and -0.5 or 0.5)))
    end
    local s = string.format("%.4f", n)
    s = s:gsub("0+$", ""):gsub("%.$", "")
    return s
end

-- ── 4. Try offline first ──────────────────────────────────────────────────

local normalised = normalise(query)
jarvis.log("info", "math normalised: " .. normalised)

local result_num =
       handle_sqrt(normalised)
    or handle_percent(normalised)
    or shunting_yard_eval(tokenize(normalised))

if result_num ~= nil then
    local pretty = format_result(result_num)
    if pretty then
        local prefix = (lang == "ru" and "Получилось " or "Result: ")
        return jarvis.cmd.ok(prefix .. pretty)
    end
end

-- ── 5. LLM fallback for word-problems / equations / conversions ───────────

local token = jarvis.system.env("GROQ_TOKEN")
if not token or token == "" then
    return jarvis.cmd.error(lang == "ru"
        and "Не понял выражение, и GROQ_TOKEN не задан."
        or  "Couldn't parse, and GROQ_TOKEN unset.")
end

local base  = jarvis.system.env("GROQ_BASE_URL")
if not base or base == "" then base = "https://api.groq.com/openai/v1" end
local model = jarvis.system.env("GROQ_MODEL")
if not model or model == "" then model = "llama-3.3-70b-versatile" end

local sys = "Ты калькулятор. Решай арифметику, простые уравнения, конверсии единиц, проценты. "
    .. "Отвечай ТОЛЬКО результатом числом или короткой строкой. "
    .. "Если запрос — не математика, ответь одним словом «нет»."

local payload = {
    model = model,
    messages = {
        { role = "system", content = sys },
        { role = "user",   content = query },
    },
    max_tokens = 64,
    temperature = 0.0,
}

jarvis.log("info", "math fallback to LLM: " .. query)
local res = jarvis.http.post_json(base .. "/chat/completions", payload, { Authorization = "Bearer " .. token })
if not res.ok then
    return jarvis.cmd.error("Ошибка API: " .. tostring(res.status))
end
local content = (res.body or ""):match('"content"%s*:%s*"(.-[^\\])"')
if not content then
    return jarvis.cmd.error(lang == "ru" and "Не смог распарсить ответ" or "Couldn't parse response")
end
content = content:gsub('\\n', '\n'):gsub('\\"', '"'):gsub('\\\\', '\\'):gsub('\\t', '\t')
content = content:gsub("^%s+", ""):gsub("%s+$", "")
if content:lower() == "нет" or content == "" then
    return jarvis.cmd.not_found(lang == "ru" and "Это не математика" or "Not math")
end

local prefix = (lang == "ru" and "Получилось " or "Result: ")
return jarvis.cmd.ok(prefix .. content)