--
-- Licensed to the Apache Software Foundation (ASF) under one or more
-- contributor license agreements.  See the NOTICE file distributed with
-- this work for additional information regarding copyright ownership.
-- The ASF licenses this file to You under the Apache License, Version 2.0
-- (the "License"); you may not use this file except in compliance with
-- the License.  You may obtain a copy of the License at
--
--     http://www.apache.org/licenses/LICENSE-2.0
--
-- Unless required by applicable law or agreed to in writing, software
-- distributed under the License is distributed on an "AS IS" BASIS,
-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-- See the License for the specific language governing permissions and
-- limitations under the License.
--
local limit_count     = require("apisix.plugins.limit-count.init")
local core            = require("apisix.core")
local config_local    = require("apisix.core.config_local")
local gq_parse        = require("graphql").parse
local limit_count_ver = require("resty.limit.count")._VERSION

local type    = type
local pairs   = pairs
local ipairs  = ipairs
local pcall   = pcall
local max     = math.max
local tonumber = tonumber

local GRAPHQL_DEFAULT_MAX_SIZE = 1048576

local plugin_name = "graphql-limit-count"
local _M = {
    version = 0.1,
    priority = 1004,
    name = plugin_name,
    schema = limit_count.schema,
}


function _M.check_schema(conf)
    return limit_count.check_schema(conf)
end


local GRAPHQL_REQ_QUERY     = "query"
local GRAPHQL_REQ_MIME_JSON = "application/json"
local GRAPHQL_REQ_MIME_GQL  = "application/graphql"


local fetch_graphql_body = {
    ["POST"] = function(ctx, max_size)
        local body, err = core.request.get_body(max_size, ctx)
        if not body then
            return nil, "failed to read graphql data, " .. (err or "request body has zero size")
        end

        return body
    end
}


local check_graphql_request = {
    ["POST"] = function(ctx, body)
        local content_type = core.request.header(ctx, "Content-Type") or ""

        if core.string.has_prefix(content_type, GRAPHQL_REQ_MIME_JSON) then
            local res, err = core.json.decode(body)
            if not res then
                return false, "invalid graphql request, " .. err
            end

            if not res[GRAPHQL_REQ_QUERY] then
                return false, "invalid graphql request, json body[" ..
                                GRAPHQL_REQ_QUERY .. "] is nil"
            end

            return true, res[GRAPHQL_REQ_QUERY]
        end

        if core.string.has_prefix(content_type, GRAPHQL_REQ_MIME_GQL) then
            return true, body
        end

        return false, "invalid graphql request, error content-type: " .. content_type
    end
}


-- Returns the maximum selection nesting depth of the GraphQL query AST.
-- Fragment spreads are expanded in place using the provided fragment map;
-- inline fragments are treated as transparent wrappers over their selections.
-- The visited table guards against fragment definition cycles.
local function node_depth(node, fragments, visited)
    if type(node) ~= "table" then
        return 0
    end

    if node.kind == "fragmentSpread" then
        local name = node.name and node.name.value
        if not name or visited[name] then
            return 0
        end
        local frag = fragments[name]
        if not frag or not frag.selectionSet then
            return 0
        end
        visited[name] = true
        local depth = node_depth(frag.selectionSet.selections, fragments, visited)
        visited[name] = nil
        return depth
    end

    if node.kind == "inlineFragment" then
        if not node.selectionSet then
            return 0
        end
        return node_depth(node.selectionSet.selections, fragments, visited)
    end

    local depth = 0
    for k, v in pairs(node) do
        local child
        if k == "selections" then
            child = 1 + node_depth(v, fragments, visited)
        else
            child = node_depth(v, fragments, visited)
        end
        depth = max(depth, child)
    end

    return depth
end


function _M.access(conf, ctx)
    if limit_count_ver < '1.0.0' then
        core.log.error("need to build APISIX-Base to support GraphQL limit count")
        return 501
    end

    local method = core.request.get_method()
    if method ~= "POST" then
        return 405
    end

    local max_size = GRAPHQL_DEFAULT_MAX_SIZE
    local local_conf = config_local.local_conf()
    if local_conf then
        local size = core.table.try_read_attr(local_conf, "graphql", "max_size")
        if size then
            local size_num = tonumber(size)
            if size_num and size_num > 0 then
                max_size = size_num
            end
        end
    end

    local body, err = fetch_graphql_body[method](ctx, max_size)
    if not body then
        core.log.error(err)
        return 400, {message = "Invalid graphql request: can't get graphql request body"}
    end

    local is_graphql_req, query_or_err = check_graphql_request[method](ctx, body)
    if not is_graphql_req then
        core.log.error(query_or_err)
        return 400, {message = query_or_err}
    end

    local ok, res = pcall(gq_parse, query_or_err)
    if not ok then
        core.log.error("failed to parse graphql: ", res)
        return 400, {message = "Invalid graphql request: failed to parse graphql query"}
    end

    -- Split definitions into executable operations and named fragment definitions.
    local fragments = {}
    local operations = {}
    for _, def in ipairs(res.definitions) do
        if def.kind == "fragmentDefinition" then
            fragments[def.name.value] = def
        else
            operations[#operations + 1] = def
        end
    end

    if #operations == 0 then
        core.log.error("failed to parse graphql: empty query")
        return 400, {message = "Invalid graphql request: empty graphql query"}
    end

    local depth = 0
    for _, op in ipairs(operations) do
        local d = node_depth(op, fragments, {})
        depth = max(depth, d)
    end
    depth = max(depth, 1)
    core.log.info("graphql query depth: ", depth)

    return limit_count.rate_limit(conf, ctx, plugin_name, depth)
end


return _M
