#include "create_chat_completion.h"

namespace humanus {

json CreateChatCompletion::_build_parameters() const {
    if (response_type == "string") {
        return {
            {"type", "object"},
            {"properties", {
                {"response", {
                    {"type", "string"},
                    {"description", "The response text that should be delivered to the user."}
                }}
            }},
            {"required", required}
        };
    }

    // TODO: handle other types

    return _create_type_schema(response_type);
}

json CreateChatCompletion::_create_type_schema(const std::string& type_hint) const {
    // 处理基本类型
    if (type_mapping.find(type_hint) != type_mapping.end()) {
        return {
            {"type", "object"},
            {"properties", {
                {"response", {
                    {"type", type_mapping.at(type_hint)},
                    {"description", "Response of type " + type_hint}
                }}
            }},
            {"required", required}
        };
    }

    // 处理数组类型
    if (type_hint.find("array") == 0) {
        std::string item_type = "string"; // 默认项类型
        return {
            {"type", "object"},
            {"properties", {
                {"response", {
                    {"type", "array"},
                    {"items", {{"type", type_mapping.at(item_type)}}}
                }}
            }},
            {"required", required}
        };
    }

    // 处理字典类型
    if (type_hint.find("object") == 0) {
        return {
            {"type", "object"},
            {"properties", {
                {"response", {
                    {"type", "object"},
                    {"additionalProperties", {{"type", "string"}}}
                }}
            }},
            {"required", required}
        };
    }

    // 默认返回字符串类型
    return {
        {"type", "object"},
        {"properties", {
            {"response", {
                {"type", "string"},
                {"description", "The response text that should be delivered to the user."}
            }}
        }},
        {"required", required}
    };
}

json CreateChatCompletion::_get_type_info(const std::string& type_hint) const {
    return {
        {"type", type_mapping.count(type_hint) ? type_mapping.at(type_hint) : "string"},
        {"description", "Value of type " + type_hint}
    };
}

json CreateChatCompletion::_create_union_schema(const std::vector<std::string>& types) const {
    json type_infos = json::array();
    for (const auto& t : types) {
        type_infos.push_back(_get_type_info(t));
    }

    return {
        {"type", "object"},
        {"properties", {
            {"response", {{"anyOf", type_infos}}}
        }},
        {"required", required}
    };
}

// Execute the chat completion with type conversion.
ToolResult CreateChatCompletion::execute(const json& args) {
    std::vector<std::string> req_fields = args.contains("required") ? args["required"].get<std::vector<std::string>>() : required;
    
    // Handle case when required is a list
    if (!req_fields.empty()) {
        if (req_fields.size() == 1) {
            std::string required_field = req_fields[0];
            return ToolResult(args.contains(required_field) ? args.at(required_field) : "");
        } else {
            // Return multiple fields as an object
            json result = json::object();
            for (const auto& field : req_fields) {
                result[field] = args.contains(field) ? args.at(field) : "";
            }
            return ToolResult(result);
        }
    } else {
        std::string required_field = "response";
        return ToolResult(args.contains(required_field) ? args.at(required_field) : "");
    }

    // TODO: handle other types (Only string and array are supported for now)
}

} // namespace humanuse