using Test
using YiemAgent
using YiemAgent.agentCore
using YiemAgent.type
using JSON

# Import the function from the private module scope
import YiemAgent.agentCore: _extractToolCalls

@testset "_extractToolCalls" begin

    # -------------------------------------------------------------- #
    # Format 1: response["message"]["tool_calls"] (LMStudio.jl style) #
    # -------------------------------------------------------------- #

    @testset "single tool call via message format" begin
        response = Dict{String,Any}(
            "finish_reason" => "tool_calls",
            "index" => 0,
            "message" => Dict{String,Any}(
                "role" => "assistant",
                "content" => "",
                "reasoning_content" => "Let me check the weather.",
                "tool_calls" => Any[
                    Dict{String,Any}(
                        "type" => "function",
                        "function" => Dict{String,Any}(
                            "name" => "getWeather",
                            "arguments" => "{\"city\":\"Bangkok, Thailand\"}",
                        ),
                        "id" => "tc_001",
                    )
                ],
            ),
        )
        has_toolcalls, tc_list = _extractToolCalls(response)
        @test has_toolcalls == true
        @test length(tc_list) == 1
        @test tc_list[1].name == "getWeather"
        @test tc_list[1].id == "tc_001"
        @test tc_list[1].arguments["city"] == "Bangkok, Thailand"
    end

    @testset "multiple tool calls via message format" begin
        response = Dict{String,Any}(
            "message" => Dict{String,Any}(
                "role" => "assistant",
                "content" => "",
                "tool_calls" => Any[
                    Dict{String,Any}(
                        "type" => "function",
                        "function" => Dict{String,Any}(
                            "name" => "getWeather",
                            "arguments" => "{\"city\":\"Tokyo, Japan\"}",
                        ),
                        "id" => "tc_001",
                    ),
                    Dict{String,Any}(
                        "type" => "function",
                        "function" => Dict{String,Any}(
                            "name" => "getTime",
                            "arguments" => "{\"timezone\":\"Asia/Tokyo\"}",
                        ),
                        "id" => "tc_002",
                    ),
                ],
            ),
        )
        has_toolcalls, tc_list = _extractToolCalls(response)
        @test has_toolcalls == true
        @test length(tc_list) == 2
        @test tc_list[1].name == "getWeather"
        @test tc_list[1].arguments["city"] == "Tokyo, Japan"
        @test tc_list[2].name == "getTime"
        @test tc_list[2].arguments["timezone"] == "Asia/Tokyo"
    end

    @testset "tool call with empty arguments string" begin
        response = Dict{String,Any}(
            "message" => Dict{String,Any}(
                "role" => "assistant",
                "tool_calls" => Any[
                    Dict{String,Any}(
                        "type" => "function",
                        "function" => Dict{String,Any}(
                            "name" => "listTools",
                            "arguments" => "{}",
                        ),
                        "id" => "tc_empty",
                    ),
                ],
            ),
        )
        has_toolcalls, tc_list = _extractToolCalls(response)
        @test has_toolcalls == true
        @test length(tc_list) == 1
        @test tc_list[1].name == "listTools"
        @test tc_list[1].arguments == Dict{String,Any}()
    end

    @testset "tool call with missing id falls back to uuid" begin
        response = Dict{String,Any}(
            "message" => Dict{String,Any}(
                "role" => "assistant",
                "tool_calls" => Any[
                    Dict{String,Any}(
                        "type" => "function",
                        "function" => Dict{String,Any}(
                            "name" => "getTime",
                            "arguments" => "{\"city\":\"NYC\"}",
                        ),
                    ),
                ],
            ),
        )
        has_toolcalls, tc_list = _extractToolCalls(response)
        @test has_toolcalls == true
        @test length(tc_list) == 1
        @test !isempty(tc_list[1].id)
        @test tc_list[1].name == "getTime"
    end

    @testset "tool call with non-string arguments (pre-parsed)" begin
        response = Dict{String,Any}(
            "message" => Dict{String,Any}(
                "role" => "assistant",
                "tool_calls" => Any[
                    Dict{String,Any}(
                        "type" => "function",
                        "function" => Dict{String,Any}(
                            "name" => "getWeather",
                            "arguments" => Dict{String,Any}("city" => "London", "units" => "fahrenheit"),
                        ),
                        "id" => "tc_parsed",
                    ),
                ],
            ),
        )
        has_toolcalls, tc_list = _extractToolCalls(response)
        @test has_toolcalls == true
        @test length(tc_list) == 1
        @test tc_list[1].arguments["city"] == "London"
        @test tc_list[1].arguments["units"] == "fahrenheit"
    end

    # ----------------------------------------------------------- #
    # Format 2: response.content blocks (OpenAI API style)        #
    # ----------------------------------------------------------- #

    @testset "content blocks with tool_calls" begin
        response = Dict{String,Any}(
            "content" => Any[
                Dict{String,Any}("type" => "text", "text" => "Let me check."),
                Dict{String,Any}(
                    "type" => "tool_calls",
                    "tool_calls" => Any[
                        Dict{String,Any}(
                            "type" => "function",
                            "function" => Dict{String,Any}(
                                "name" => "getWeather",
                                "arguments" => "{\"city\":\"Paris\"}",
                            ),
                            "id" => "tc_block_1",
                        ),
                    ],
                ),
            ],
        )
        has_toolcalls, tc_list = _extractToolCalls(response)
        @test has_toolcalls == true
        @test length(tc_list) == 1
        @test tc_list[1].name == "getWeather"
        @test tc_list[1].arguments["city"] == "Paris"
    end

    @testset "content blocks with tool_call (single-call format)" begin
        response = Dict{String,Any}(
            "content" => Any[
                Dict{String,Any}(
                    "type" => "tool_call",
                    "id" => "tc_single",
                    "name" => "getTime",
                    "arguments" => Dict{String,Any}("timezone" => "Europe/London"),
                ),
            ],
        )
        has_toolcalls, tc_list = _extractToolCalls(response)
        @test has_toolcalls == true
        @test length(tc_list) == 1
        @test tc_list[1].name == "getTime"
        @test tc_list[1].id == "tc_single"
        @test tc_list[1].arguments["timezone"] == "Europe/London"
    end

    # ----------------------------------------------------------- #
    # Format 2 via struct-like object (no .content field)         #
    # ----------------------------------------------------------- #

    @testset "no tool calls found" begin
        response = Dict{String,Any}(
            "content" => Any[
                Dict{String,Any}("type" => "text", "text" => "Hello world."),
            ],
        )
        has_toolcalls, tc_list = _extractToolCalls(response)
        @test has_toolcalls == false
        @test length(tc_list) == 0
    end

    @testset "empty message" begin
        response = Dict{String,Any}()
        has_toolcalls, tc_list = _extractToolCalls(response)
        @test has_toolcalls == false
        @test length(tc_list) == 0
    end

    @testset "message with empty tool_calls array" begin
        response = Dict{String,Any}(
            "message" => Dict{String,Any}(
                "role" => "assistant",
                "tool_calls" => Any[],
            ),
        )
        has_toolcalls, tc_list = _extractToolCalls(response)
        @test has_toolcalls == false
        @test length(tc_list) == 0
    end

    @testset "Format 1 takes priority over Format 2" begin
        response = Dict{String,Any}(
            "message" => Dict{String,Any}(
                "role" => "assistant",
                "tool_calls" => Any[
                    Dict{String,Any}(
                        "type" => "function",
                        "function" => Dict{String,Any}(
                            "name" => "getWeather",
                            "arguments" => "{\"city\":\"Format1\"}",
                        ),
                        "id" => "tc_fmt1",
                    ),
                ],
            ),
            "content" => Any[
                Dict{String,Any}(
                    "type" => "tool_call",
                    "id" => "tc_fmt2",
                    "name" => "getTime",
                    "arguments" => Dict{String,Any}("city" => "Format2"),
                ),
            ],
        )
        has_toolcalls, tc_list = _extractToolCalls(response)
        @test has_toolcalls == true
        @test length(tc_list) == 1
        @test tc_list[1].name == "getWeather"
    end

    # ----------------------------------------------------------- #
    # edge cases                                                    #
    # ----------------------------------------------------------- #

    @testset "tool call with null arguments" begin
        response = Dict{String,Any}(
            "message" => Dict{String,Any}(
                "role" => "assistant",
                "tool_calls" => Any[
                    Dict{String,Any}(
                        "type" => "function",
                        "function" => Dict{String,Any}(
                            "name" => "getTime",
                            "arguments" => nothing,
                        ),
                        "id" => "tc_null",
                    ),
                ],
            ),
        )
        has_toolcalls, tc_list = _extractToolCalls(response)
        @test has_toolcalls == true
        @test length(tc_list) == 1
        @test tc_list[1].name == "getTime"
    end

    @testset "tool call with missing function key" begin
        response = Dict{String,Any}(
            "message" => Dict{String,Any}(
                "role" => "assistant",
                "tool_calls" => Any[
                    Dict{String,Any}(
                        "type" => "function",
                        "id" => "tc_nofunc",
                    ),
                ],
            ),
        )
        has_toolcalls, tc_list = _extractToolCalls(response)
        @test has_toolcalls == true
        @test length(tc_list) == 1
        @test tc_list[1].name == ""
    end

    @testset "tool call with missing name in function block" begin
        response = Dict{String,Any}(
            "message" => Dict{String,Any}(
                "role" => "assistant",
                "tool_calls" => Any[
                    Dict{String,Any}(
                        "type" => "function",
                        "function" => Dict{String,Any}("arguments" => "{}"),
                        "id" => "tc_noname",
                    ),
                ],
            ),
        )
        has_toolcalls, tc_list = _extractToolCalls(response)
        @test has_toolcalls == true
        @test length(tc_list) == 1
        @test tc_list[1].name == ""
    end

    @testset "message format with JSON.Object (JSON.parse result)" begin
        json_str = JSON.json(Dict(
            "message" => Dict(
                "role" => "assistant",
                "tool_calls" => [
                    Dict(
                        "type" => "function",
                        "function" => Dict("name" => "getWeather", "arguments" => "{\"city\":\"Test\"}"),
                        "id" => "tc_jsonobj",
                    ),
                ],
            ),
        ))
        parsed = JSON.parse(json_str)
        has_toolcalls, tc_list = _extractToolCalls(parsed)
        @test has_toolcalls == true
        @test length(tc_list) == 1
        @test tc_list[1].name == "getWeather"
        @test tc_list[1].arguments["city"] == "Test"
    end

end