--
-- 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 core = require("apisix.core")
local secret = require("apisix.secret")
local schema = require("apisix.plugins.ai-proxy.schema")
local base   = require("apisix.plugins.ai-proxy.base")
local plugin = require("apisix.plugin")
local ipmatcher  = require("resty.ipmatcher")
local healthcheck_manager = require("apisix.healthcheck_manager")
local resource = require("apisix.resource")
local exporter = require("apisix.plugins.prometheus.exporter")
local tonumber = tonumber
local pairs = pairs
local table_sort = table.sort
local math_random = math.random
local ngx_now = ngx.now

local require = require
local pcall = pcall
local ipairs = ipairs
local type = type
local string = string
local url = require("socket.url")

local priority_balancer = require("apisix.balancer.priority")
local endpoint_regex = "^(https?)://([^:/]+):?(%d*)/?.*$"

local pickers = {}
local lrucache_server_picker = core.lrucache.new({
    ttl = 300, count = 256
})

local plugin_name = "ai-proxy-multi"
local _M = {
    version = 0.5,
    priority = 1041,
    name = plugin_name,
    schema = schema.ai_proxy_multi_schema,
}

local function fallback_strategy_has(strategy, name)
    if not strategy then
        return false
    end

    if type(strategy) == "string" then
        return strategy == name
    end

    if type(strategy) == "table" then
        for _, v in ipairs(strategy) do
            if v == name then
                return true
            end
        end
    end

    return false
end


local function get_chash_key_schema(hash_on)
    if hash_on == "vars" then
        return core.schema.upstream_hash_vars_schema
    end

    if hash_on == "header" or hash_on == "cookie" then
        return core.schema.upstream_hash_header_schema
    end

    if hash_on == "consumer" then
        return nil, nil
    end

    if hash_on == "vars_combinations" then
        return core.schema.upstream_hash_vars_combinations_schema
    end

    return nil, "invalid hash_on type " .. hash_on
end


function _M.check_schema(conf)
    local ok, err = core.schema.check(schema.ai_proxy_multi_schema, conf)
    if not ok then
        return false, err
    end

    for _, instance in ipairs(conf.instances) do
        local endpoint = instance and instance.override and instance.override.endpoint
        if endpoint then
            local scheme, host, _ = endpoint:match(endpoint_regex)
            if not scheme or not host  then
                return false, "invalid endpoint"
            end
        end
        local ai_provider, err = pcall(require, "apisix.plugins.ai-providers." .. instance.provider)
        if not ai_provider then
            core.log.warn("fail to require ai provider: ", instance.provider, ", err", err)
            return false, "ai provider: " .. instance.provider .. " is not supported."
        end
        local sa_json = core.table.try_read_attr(instance, "auth", "gcp", "service_account_json")
        if sa_json and not secret.is_secret_ref(sa_json) then
            local _, err = core.json.decode(sa_json)
            if err then
                return false, "invalid gcp service_account_json: " .. err
            end
        end
        local ok, err = schema.validate_provider_requirements(instance)
        if not ok then
            return false, "instance '" .. (instance.name or "?") .. "': " .. err
        end
    end
    local algo = core.table.try_read_attr(conf, "balancer", "algorithm")
    local hash_on = core.table.try_read_attr(conf, "balancer", "hash_on")
    local hash_key = core.table.try_read_attr(conf, "balancer", "key")

    if type(algo) == "string" and algo == "chash" then
        if not hash_on then
            return false, "must configure `hash_on` when balancer algorithm is chash"
        end

        if hash_on ~= "consumer" and not hash_key then
            return false, "must configure `hash_key` when balancer `hash_on` is not set to cookie"
        end

        local key_schema, err = get_chash_key_schema(hash_on)
        if err then
            return false, "type is chash, err: " .. err
        end

        if key_schema then
            local ok, err = core.schema.check(key_schema, hash_key)
            if not ok then
                return false, "invalid configuration: " .. err
            end
        end
    end

    return ok
end


local function transform_instances(new_instances, instance)
    if not new_instances._priority_index then
        new_instances._priority_index = {}
    end

    if not new_instances[instance.priority] then
        new_instances[instance.priority] = {}
        core.table.insert(new_instances._priority_index, instance.priority)
    end

    new_instances[instance.priority][instance.name] = instance.weight
end

local function sort_nodes(a, b)
    if a.host == b.host then
        return (a.port or 0) < (b.port or 0)
    end
    return a.host < b.host
end


local function nodes_equal(old_nodes, new_nodes)
    if old_nodes == new_nodes then
        return true
    end

    if type(old_nodes) ~= "table" or #old_nodes ~= #new_nodes then
        return false
    end

    for i, new_node in ipairs(new_nodes) do
        local old_node = old_nodes[i]
        for _, field in ipairs({"host", "port", "scheme", "domain"}) do
            if old_node[field] ~= new_node[field] then
                return false
            end
        end
    end

    return true
end


local function parse_domain_for_nodes(node)
    local host = node.domain or node.host
    if not ipmatcher.parse_ipv4(host)
       and not ipmatcher.parse_ipv6(host)
    then
        local ips, err = core.resolver.parse_domain_all(host)
        if err then
            core.log.error("dns resolver domain: ", host, " error: ", err)
        end

        if ips then
            local nodes = core.table.new(#ips, 0)
            for _, ip in ipairs(ips) do
                local new_node = core.table.clone(node)
                new_node.host = ip
                new_node.domain = host
                core.table.insert(nodes, new_node)
            end
            table_sort(nodes, sort_nodes)
            return nodes
        end
    end

    return {node}
end


local function make_endpoint(node)
    local host = node.host
    if ipmatcher.parse_ipv6(host) then
        host = "[" .. host .. "]"
    end

    local endpoint = node.scheme .. "://" .. host .. ":" .. node.port
    if node.path then
        endpoint = endpoint .. node.path
    end
    if node.query then
        endpoint = endpoint .. "?" .. node.query
    end
    return endpoint
end


local function make_host_header(node)
    if not node.domain then
        return nil
    end

    local port = tonumber(node.port)
    if (node.scheme == "https" and port ~= 443)
       or (node.scheme ~= "https" and port ~= 80)
    then
        return node.domain .. ":" .. node.port
    end

    return node.domain
end


local function use_node_for_request(instance_conf, node)
    if not node then
        return
    end

    instance_conf._dns_value = node
    instance_conf._resolved_endpoint = make_endpoint(node)
    instance_conf._resolved_host_header = make_host_header(node)
    instance_conf._resolved_ssl_server_name = node.domain
end


local function pick_request_node(nodes)
    if not nodes or #nodes == 0 then
        return
    end

    return nodes[math_random(1, #nodes)]
end


-- resolves endpoint and sets it on _dns_nodes
local function resolve_endpoint(instance_conf)
    local scheme, host, port, path, query
    local endpoint = core.table.try_read_attr(instance_conf, "override", "endpoint")
    if endpoint then
        local parsed = url.parse(endpoint)
        scheme = parsed.scheme
        host = parsed.host
        port = parsed.port
        path = parsed.path
        query = parsed.query
        if not port then
            port = (scheme == "https") and 443 or 80
        end
        port = tonumber(port)
    else
        local ai_provider = require("apisix.plugins.ai-providers." .. instance_conf.provider)
        if ai_provider.get_node then
            local node = ai_provider.get_node(instance_conf)
            host = node.host
            port = node.port
        else
            host = ai_provider.host
            port = ai_provider.port
        end
        scheme = "https"
    end

    local new_node = {
        host = host,
        port = port,
        scheme = scheme,
        path = path,
        query = query,
    }
    local new_nodes = parse_domain_for_nodes(new_node)

    local nodes_changed = not nodes_equal(instance_conf._dns_nodes, new_nodes)

    if nodes_changed then
        instance_conf._dns_nodes = new_nodes
        instance_conf._nodes_ver = (instance_conf._nodes_ver or 0) + 1
        core.log.info("DNS resolution changed for instance: ", instance_conf.name,
                     " new nodes: ", core.json.delay_encode(new_nodes))
    end

    use_node_for_request(instance_conf, pick_request_node(instance_conf._dns_nodes))
end


local function get_checkers_status_ver(checkers)
    local status_ver_total = 0
    for _, checker in pairs(checkers) do
        status_ver_total = status_ver_total + checker.status_ver
    end
    return status_ver_total
end


local function fetch_health_instances(conf, checkers)
    local instances = conf.instances
    local new_instances = core.table.new(0, #instances)
    if not checkers then
        for _, ins in ipairs(conf.instances) do
            transform_instances(new_instances, ins)
        end
        return new_instances
    end

    for _, ins in ipairs(instances) do
        local checker = checkers[ins.name]
        if checker then
            local host = ins.checks and ins.checks.active and ins.checks.active.host
            local port = ins.checks and ins.checks.active and ins.checks.active.port
            local healthy_nodes = {}
            ins._healthy_dns_nodes = nil

            for _, node in ipairs(ins._dns_nodes or {}) do
                local ok, err = checker:get_target_status(node.host, port or node.port, host)
                if ok then
                    healthy_nodes[#healthy_nodes + 1] = node
                elseif err then
                    core.log.warn("failed to get health check target status, addr: ",
                        node.host, ":", port or node.port, ", host: ", host, ", err: ", err)
                end
            end

            if #healthy_nodes > 0 then
                ins._healthy_dns_nodes = healthy_nodes
                transform_instances(new_instances, ins)
            end
        else
            ins._healthy_dns_nodes = nil
            transform_instances(new_instances, ins)
        end
    end

    if core.table.nkeys(new_instances) == 0 then
        core.log.warn("all upstream nodes is unhealthy, use default")
        for _, ins in ipairs(instances) do
            transform_instances(new_instances, ins)
        end
    end

    return new_instances
end


local function create_server_picker(conf, ups_tab, checkers)
    local picker = pickers[conf.balancer.algorithm] -- nil check
    if not picker then
        pickers[conf.balancer.algorithm] = require("apisix.balancer." .. conf.balancer.algorithm)
        picker = pickers[conf.balancer.algorithm]
    end

    local new_instances = fetch_health_instances(conf, checkers)
    core.log.info("fetch health instances: ", core.json.delay_encode(new_instances))

    if #new_instances._priority_index > 1 then
        core.log.info("new instances: ", core.json.delay_encode(new_instances))
        return priority_balancer.new(new_instances, ups_tab, picker)
    end
    core.log.info("upstream nodes: ",
                core.json.delay_encode(new_instances[new_instances._priority_index[1]]))
    return picker.new(new_instances[new_instances._priority_index[1]], ups_tab)
end


local function get_instance_conf(instances, name)
    for _, ins in ipairs(instances) do
        if ins.name == name then
            return ins
        end
    end
end


local function pick_target(ctx, conf, ups_tab)
    local checkers = {}
    local res_conf = resource.fetch_latest_conf(conf._meta.parent.resource_key)
    if not res_conf then
        return nil, nil, "failed to fetch the parent config"
    end
    local instances = res_conf.value.plugins[plugin_name].instances
    for i, instance in ipairs(conf.instances) do
        if instance.checks then
            resolve_endpoint(instance)
            -- json path is 0 indexed so we need to decrement i
            local resource_path = conf._meta.parent.resource_key ..
                                  "#plugins['ai-proxy-multi'].instances[" .. i-1 .. "]"
            local resource_version = conf._meta.parent.resource_version
            if instance._nodes_ver then
                resource_version = resource_version .. instance._nodes_ver
            end
            instances[i]._dns_nodes = instance._dns_nodes
            instances[i]._nodes_ver = instance._nodes_ver
            local checker = healthcheck_manager.fetch_checker(resource_path, resource_version)
            checkers[instance.name] = checker
        end
    end

    local version = plugin.conf_version(conf)
    if checkers then
        local status_ver = get_checkers_status_ver(checkers)
        version = version .. "#" .. status_ver
    end

    local server_picker = ctx.server_picker
    if not server_picker then
        server_picker = lrucache_server_picker(ctx.matched_route.key, version,
                                               create_server_picker, conf, ups_tab, checkers)
    end
    if not server_picker then
        return nil, nil, "failed to fetch server picker"
    end
    ctx.server_picker = server_picker

    local instance_name, err = server_picker.get(ctx)
    if err then
        return nil, nil, err
    end
    ctx.balancer_server = instance_name
    if conf.fallback_strategy == "instance_health_and_rate_limiting" or -- for backwards compatible
       fallback_strategy_has(conf.fallback_strategy, "rate_limiting") then
        local ai_rate_limiting = require("apisix.plugins.ai-rate-limiting")
        for _ = 1, #conf.instances do
            if ai_rate_limiting.check_instance_status(nil, ctx, instance_name) then
                break
            end
            core.log.warn("ai instance: ", instance_name,
                             " is not available, try to pick another one")
            server_picker.after_balance(ctx, true)
            instance_name, err = server_picker.get(ctx)
            if err then
                return nil, nil, err
            end
            ctx.balancer_server = instance_name
        end
    end

    local instance_conf = get_instance_conf(conf.instances, instance_name)
    local nodes = instance_conf._healthy_dns_nodes or instance_conf._dns_nodes
    use_node_for_request(instance_conf, pick_request_node(nodes))
    return instance_name, instance_conf
end


local function pick_ai_instance(ctx, conf, ups_tab)
    local instance_name, instance_conf, err
    if #conf.instances == 1 then
        instance_name = conf.instances[1].name
        instance_conf = conf.instances[1]
    else
        instance_name, instance_conf, err = pick_target(ctx, conf, ups_tab)
    end

    core.log.info("picked instance: ", instance_name)
    return instance_name, instance_conf, err
end

function _M.access(conf, ctx)
    -- Detect the client protocol and read the body first. get_json_request_body_table
    -- reads and size-checks the body exactly once (bounded by max_req_body_size,
    -- rejecting via Content-Length before buffering), so oversized requests are
    -- rejected before any balancer / DNS / health-check work below.
    local err, code = base.detect_request_type(ctx, conf.max_req_body_size)
    if err then
        return code or 400, err
    end

    local ups_tab = {}
    local algo = core.table.try_read_attr(conf, "balancer", "algorithm")
    if algo == "chash" then
        local hash_on = core.table.try_read_attr(conf, "balancer", "hash_on")
        local hash_key = core.table.try_read_attr(conf, "balancer", "key")
        ups_tab["key"] = hash_key
        ups_tab["hash_on"] = hash_on
    end

    local name, ai_instance, perr = pick_ai_instance(ctx, conf, ups_tab)
    if perr then
        return 503, perr
    end
    ctx.picked_ai_instance_name = name
    ctx.picked_ai_instance = ai_instance
    ctx.balancer_ip = name
    ctx.bypass_nginx_upstream = true
end


local function retry_on_error(ctx, conf, code)
    if not ctx.server_picker then
        return code
    end
    ctx.server_picker.after_balance(ctx, true)
    if (code == 429 and fallback_strategy_has(conf.fallback_strategy, "http_429")) or
       (code >= 500 and code < 600 and
       fallback_strategy_has(conf.fallback_strategy, "http_5xx")) then
        -- Slow-failure guard: only retry when the failed attempt finished within
        -- retry_on_failure_within_ms. A slow failure (e.g. a 5xx returned after
        -- minutes) is given back to the client directly, so fallback never doubles
        -- the client's wait time. ctx.llm_request_start_time is reset by base
        -- before_proxy at the start of every attempt, so this measures the elapsed
        -- time of the attempt that just failed.
        if conf.retry_on_failure_within_ms and ctx.llm_request_start_time then
            local elapsed_ms = (ngx_now() - ctx.llm_request_start_time) * 1000
            if elapsed_ms > conf.retry_on_failure_within_ms then
                core.log.warn("ai instance failed after ", elapsed_ms,
                              "ms, exceeding retry_on_failure_within_ms ",
                              conf.retry_on_failure_within_ms, ", not retrying")
                return code
            end
        end

        -- Cap the number of fallback retries so a single request does not exhaust
        -- every instance when many are configured.
        if conf.max_retries then
            ctx.ai_retries = (ctx.ai_retries or 0) + 1
            if ctx.ai_retries > conf.max_retries then
                core.log.warn("reached max_retries ", conf.max_retries,
                              ", not retrying")
                return code
            end
        end

        local name, ai_instance, err = pick_ai_instance(ctx, conf)
        if err then
            core.log.error("failed to pick new AI instance: ", err)
            return 502
        end
        ctx.balancer_ip = name
        ctx.picked_ai_instance_name = name
        ctx.picked_ai_instance = ai_instance
        return
    end
    return code
end

function _M.construct_upstream(instance)
    if not instance then
        return nil, "instance configuration is nil"
    end
    local upstream = {}
    local nodes = instance._dns_nodes
    if not nodes then
        resolve_endpoint(instance)
        nodes = instance._dns_nodes
        if not nodes then
            return nil, "failed to resolve endpoint for instance: " .. instance.name
        end
    end

    local upstream_nodes = core.table.new(#nodes, 0)
    for _, node in ipairs(nodes) do
        if not node.host or not node.port then
            return nil, "invalid upstream node: missing host or port"
        end

        core.table.insert(upstream_nodes, {
            host = node.host,
            port = node.port,
            weight = 1,
            priority = 0,
            domain = node.domain,
        })
    end

    local checks = instance.checks
    local auth = instance.auth or {}
    if checks and checks.active then
        -- Clone checks to avoid in-place mutation across requests
        checks = core.table.deepcopy(checks)
        if auth.header then
            local add_headers = {}
            checks.active.req_headers = checks.active.req_headers or {}
            for _, v in ipairs(checks.active.req_headers) do
                add_headers[v] = true
            end
            for k, v in pairs(auth.header) do
                local header = string.format("%s: %s", k, v)
                if not add_headers[header] then
                    core.table.insert(checks.active.req_headers, header)
                end
            end
        end
        if auth.query then
            checks.active.http_path = string.format("%s?%s",
                    checks.active.http_path, core.string.encode_args(auth.query))
        end
    end
    upstream.nodes = upstream_nodes
    upstream.checks = checks
    upstream._nodes_ver = instance._nodes_ver
    return upstream
end


function _M.before_proxy(conf, ctx)
     return base.before_proxy(conf, ctx, function (ctx, conf, code)
        return retry_on_error(ctx, conf, code)
    end)
end

function _M.log(conf, ctx)
    if ctx.llm_active_connections_tracked then
        exporter.dec_llm_active_connections(ctx)
        ctx.llm_active_connections_tracked = false
    end
    if conf.logging then
        base.set_logging(ctx, conf.logging.summaries, conf.logging.payloads)
    end
end

return _M
