--
-- 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.
--

--- OpenAI Responses API protocol adapter.
-- Handles the Responses API format with its different SSE event model
-- and response structure (output[] instead of choices[]).

local core = require("apisix.core")
local uuid = require("resty.jit-uuid")
local sse = require("apisix.plugins.ai-transport.sse")
local type = type
local ipairs = ipairs
local table = table
local string_sub = string.sub

local _M = {}


--- Detect whether the request matches OpenAI Responses API format.
-- Requires URI suffix (/v1/responses) and body (has input field).
function _M.matches(body, ctx)
    local uri = ctx.var and ctx.var.uri
    return uri and string_sub(uri, -13) == "/v1/responses"
        and type(body) == "table" and body.input ~= nil
end


--- Check whether the request is a streaming request.
function _M.is_streaming(body)
    return body.stream == true
end



function _M.parse_sse_event(event, ctx, state)
    if event.type == "response.output_text.delta" then
        local data, err = core.json.decode(event.data, { null_as_nil = true })
        if not data then
            core.log.warn("failed to decode SSE data: ", err)
            return { type = "skip" }
        end
        if type(data.delta) == "string" then
            return {
                type = "delta",
                texts = { data.delta },
            }
        end
        return { type = "skip" }

    elseif event.type == "response.completed" then
        local result = { type = "done" }
        local data, err = core.json.decode(event.data, { null_as_nil = true })
        if not data then
            core.log.warn("failed to decode response.completed SSE data: ", err)
            return result
        end
        if type(data.response) == "table" then
            local resp = data.response
            if type(resp.usage) == "table" then
                local usage = resp.usage
                result.type = "usage_and_done"
                result.usage = {
                    prompt_tokens = usage.input_tokens or 0,
                    completion_tokens = usage.output_tokens or 0,
                    total_tokens = usage.total_tokens or 0,
                    cache_read_input_tokens = type(usage.input_tokens_details) == "table"
                        and usage.input_tokens_details.cached_tokens or 0,
                    reasoning_tokens = type(usage.output_tokens_details) == "table"
                        and usage.output_tokens_details.reasoning_tokens or 0,
                }
                result.raw_usage = usage
            end
            if type(resp.output) == "table" then
                for _, item in ipairs(resp.output) do
                    if type(item) == "table" and item.type == "function_call" then
                        result.has_tool_call = true
                        break
                    end
                end
            end
        end
        return result

    elseif event.type == "response.failed"
            or event.type == "response.incomplete"
            or event.type == "error" then
        return { type = "done" }
    end

    -- All other Responses API events are silently passed through
    return { type = "skip" }
end


function _M.extract_response_text(res_body)
    if type(res_body.output) ~= "table" then
        return nil
    end
    local texts = {}
    for _, item in ipairs(res_body.output) do
        if type(item) == "table" and item.type == "message"
                and type(item.content) == "table" then
            for _, part in ipairs(item.content) do
                if part.type == "output_text" and type(part.text) == "string" then
                    core.table.insert(texts, part.text)
                end
            end
        end
    end
    if #texts > 0 then
        return table.concat(texts, " ")
    end
    return nil
end


--- Build a non-streaming request from system prompt and user content.
function _M.build_simple_request(system_prompt, user_content, opts)
    local body = {
        input = user_content,
        stream = false,
    }
    if system_prompt then
        body.instructions = system_prompt
    end
    if opts and opts.model then
        body.model = opts.model
    end
    return body
end


function _M.extract_usage(res_body)
    if type(res_body) ~= "table" or type(res_body.usage) ~= "table" then
        return nil, nil
    end
    local raw = res_body.usage
    local idetails = type(raw.input_tokens_details) == "table" and raw.input_tokens_details
    local odetails = type(raw.output_tokens_details) == "table" and raw.output_tokens_details
    local prompt = raw.input_tokens or 0
    local completion = raw.output_tokens or 0
    return {
        prompt_tokens = prompt,
        completion_tokens = completion,
        total_tokens = raw.total_tokens or (prompt + completion),
        cache_read_input_tokens = idetails and idetails.cached_tokens or 0,
        cache_creation_input_tokens = 0,
        reasoning_tokens = odetails and odetails.reasoning_tokens or 0,
    }, raw
end


--- Detect whether a non-streaming response contains tool calls.
function _M.has_tool_call(res_body)
    if type(res_body) ~= "table" or type(res_body.output) ~= "table" then
        return false
    end
    for _, item in ipairs(res_body.output) do
        if type(item) == "table" and item.type == "function_call" then
            return true
        end
    end
    return false
end


--- Extract the end-user identifier from a request body.
function _M.extract_end_user_id(body)
    if type(body) ~= "table" then
        return nil
    end
    if type(body.safety_identifier) == "string" then
        return body.safety_identifier
    end
    if type(body.user) == "string" then
        return body.user
    end
    return nil
end


-- Append an input item's text into `contents`. An item may be a plain string, a
-- message object ({content = string | {parts}}), or a bare content part
-- ({text = "..."}); text parts are flattened so consumers get plain strings.
local function append_item_text(contents, item)
    if type(item) == "string" then
        core.table.insert(contents, item)
        return
    end
    if type(item) ~= "table" then
        return
    end
    local content = item.content
    if type(content) == "string" then
        core.table.insert(contents, content)
    elseif type(content) == "table" then
        for _, part in ipairs(content) do
            if type(part) == "table" and type(part.text) == "string" then
                core.table.insert(contents, part.text)
            end
        end
    elseif content == nil and type(item.text) == "string" then
        core.table.insert(contents, item.text)
    elseif content == nil and type(item.output) == "string" then
        -- Responses-API tool result (function_call_output)
        core.table.insert(contents, item.output)
    elseif content == nil and type(item.output) == "table" then
        -- function_call_output.output may be an array of content parts
        for _, part in ipairs(item.output) do
            if type(part) == "table" and type(part.text) == "string" then
                core.table.insert(contents, part.text)
            end
        end
    end
end


--- Extract all text content from a request body for moderation.
function _M.extract_request_content(body)
    local contents = {}
    local input = body.input
    if type(input) == "string" then
        core.table.insert(contents, input)
    elseif type(input) == "table" then
        for _, item in ipairs(input) do
            append_item_text(contents, item)
        end
    end
    if type(body.instructions) == "string" then
        core.table.insert(contents, body.instructions)
    end
    return contents
end


-- Roles carrying the system prompt. `developer` is what OpenAI renamed `system`
-- to on o1 and later models; both land in the same prompt slot, so the two are
-- extracted together and the `system` role selector covers both.
local SYSTEM_ROLES = {
    system = true,
    developer = true,
}


-- Whether an input item belongs to a selected turn role. A bare string is user
-- text; a role item matches when its role is in `roles`; a Responses-API tool
-- result (`function_call_output`, which has no role) matches when tool is selected.
local function turn_item_matches(item, roles)
    if type(item) == "string" then
        return roles.user and true or false
    end
    if type(item) ~= "table" then
        return false
    end
    if item.role ~= nil then
        return roles[item.role] and true or false
    end
    if roles.tool and item.type == "function_call_output" then
        return true
    end
    return false
end


-- Extract turn-role (user/tool) input text for request moderation. A plain-string
-- `input` is user content. For an `input` array, mode "last" = the latest
-- consecutive block of selected-role items, "all" = every selected-role item.
-- `instructions` (the system prompt) is handled by extract_system_content.
function _M.extract_turn_content(body, mode, roles)
    local contents = {}
    local input = body.input
    if type(input) == "string" then
        if roles.user then
            core.table.insert(contents, input)
        end
        return contents
    end
    if type(input) ~= "table" then
        return contents
    end
    local start_idx = 1
    if mode ~= "all" then
        start_idx = nil
        for i = #input, 1, -1 do
            if turn_item_matches(input[i], roles) then
                start_idx = i
            else
                break
            end
        end
        if not start_idx then
            return contents
        end
    end
    for i = start_idx, #input do
        if turn_item_matches(input[i], roles) then
            append_item_text(contents, input[i])
        end
    end
    return contents
end


-- Extract system-role text (see SYSTEM_ROLES) for request moderation. Responses
-- API carries the system prompt in `instructions`; an `input` array may also
-- hold system items.
function _M.extract_system_content(body)
    local contents = {}
    if type(body.instructions) == "string" then
        core.table.insert(contents, body.instructions)
    end
    if type(body.input) == "table" then
        for _, item in ipairs(body.input) do
            if type(item) == "table" and SYSTEM_ROLES[item.role] then
                append_item_text(contents, item)
            end
        end
    end
    return contents
end


--- Get messages in canonical {role, content} format.
-- Converts instructions + input into messages-style list.
function _M.get_messages(body)
    local messages = {}
    if type(body.instructions) == "string" then
        core.table.insert(messages, {role = "system", content = body.instructions})
    end
    local input = body.input
    if type(input) == "string" then
        core.table.insert(messages, {role = "user", content = input})
    elseif type(input) == "table" then
        for _, item in ipairs(input) do
            local role = "user"
            if type(item) == "table" and type(item.role) == "string" then
                role = item.role
            end
            local texts = {}
            append_item_text(texts, item)
            if #texts > 0 then
                core.table.insert(messages, {
                    role = role,
                    content = table.concat(texts, " "),
                })
            end
        end
    end
    return messages
end


--- Prepend messages to the request body.
-- System messages go to instructions; user messages prepend to input.
function _M.prepend_messages(body, msgs)
    if not msgs or #msgs == 0 then return end
    local parts = {}
    for _, msg in ipairs(msgs) do
        core.table.insert(parts, msg.content)
    end
    local prepend_text = table.concat(parts, "\n")
    if type(body.instructions) == "string" then
        body.instructions = prepend_text .. "\n" .. body.instructions
    else
        body.instructions = prepend_text
    end
end


--- Append messages to the request body.
function _M.append_messages(body, msgs)
    if not msgs or #msgs == 0 then return end
    local parts = {}
    for _, msg in ipairs(msgs) do
        core.table.insert(parts, msg.content)
    end
    local append_text = table.concat(parts, "\n")
    if type(body.input) == "string" then
        body.input = body.input .. "\n" .. append_text
    elseif type(body.input) == "table" then
        core.table.insert(body.input, {
            type = "message",
            role = "user",
            content = append_text,
        })
    else
        body.input = append_text
    end
end


--- Get raw request content for logging.
function _M.get_request_content(body)
    return body.input
end
-- opts: {text, model, usage, stream}
function _M.build_deny_response(opts)
    local response_obj = {
        id = uuid.generate_v4(),
        object = "response",
        status = "completed",
        model = opts.model,
        output = {{
            type = "message",
            role = "assistant",
            content = {{
                type = "output_text",
                text = opts.text,
            }},
        }},
        usage = opts.usage,
    }
    if opts.stream then
        local delta_event = {
            type = "response.output_text.delta",
            delta = opts.text,
        }
        local completed_event = {
            type = "response.completed",
            response = response_obj,
        }
        return "event: response.output_text.delta\n"
            .. "data: " .. core.json.encode(delta_event) .. "\n\n"
            .. "event: response.completed\n"
            .. "data: " .. core.json.encode(completed_event) .. "\n\n"
    else
        return core.json.encode(response_obj)
    end
end


--- Build an empty usage object.
function _M.empty_usage()
    return { input_tokens = 0, output_tokens = 0, total_tokens = 0 }
end


--- Check if an SSE event is a data event.
function _M.is_data_event(event)
    return event.type == "response.completed"
end


--- Check if an SSE event is the terminal/done event.
function _M.is_done_event(event)
    return event.type == "response.completed"
end


--- Build a terminal SSE event string.
function _M.build_done_event()
    return sse.encode({
        type = "response.completed",
        data = core.json.encode({
            type = "response.completed",
            response = { status = "completed", output = {} }
        })
    })
end


return _M
