#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}
        };
    } else if (response_type == "object") {
        // 处理对象类型
        return {
            {"type", "object"},
            {"properties", json::object()},
            {"required", required}
        };
    } else if (response_type == "array") {
        // 处理数组类型
        return {
            {"type", "object"},
            {"properties", {
                {"response", {
                    {"type", "array"},
                    {"items", json::object()}
                }}
            }},
            {"required", required}
        };
    } else {
        throw std::runtime_error("Invalid response type: " + 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 _build_parameters();
}

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}
    };
}

ToolResult CreateChatCompletion::execute(const json& args) {
    std::vector<std::string> req_fields = args.contains("req") ? args["req"].get<std::vector<std::string>>() : required;
    
    if (!req_fields.empty()) {
        if (req_fields.size() == 1) {
            std::string required_field = req_fields[0];
            return args.contains(required_field) ? args.at(required_field) : "";
        } else {
            // 返回多个字段作为对象
            json result = json::object();
            for (const auto& field : req_fields) {
                result[field] = args.contains(field) ? args.at(field) : "";
            }
            return result;
        }
    } else {
        std::string required_field = "response";
        return args.contains(required_field) ? args.at(required_field) : "";
    }
}

} // namespace humanuse