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, assistant_msg = _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].type == "function"
        @test tc_list[1].arguments["city"] == "Bangkok, Thailand"
        @test assistant_msg isa assistantMessage
        @test assistant_msg.role == "assistant"
        @test assistant_msg.stopReason == "tool_calls"
        @test length(assistant_msg.content) == 1
        @test assistant_msg.content[1] isa textContent
        @test assistant_msg.content[1].text == "Let me check the weather."
    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, assistant_msg = _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"
        @test assistant_msg.role == "assistant"
    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, assistant_msg = _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}()
        @test assistant_msg.stopReason == "end_turn"
    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, assistant_msg = _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 dict)" 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, assistant_msg = _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

    @testset "tool call with api/provider/model/usage metadata" begin
        response = Dict{String,Any}(
            "api" => "openai",
            "provider" => "anthropic",
            "model" => "claude-3-opus",
            "message" => Dict{String,Any}(
                "role" => "assistant",
                "tool_calls" => Any[
                    Dict{String,Any}(
                        "type" => "function",
                        "function" => Dict{String,Any}(
                            "name" => "getTime",
                            "arguments" => "{}",
                        ),
                        "id" => "tc_meta",
                    ),
                ],
            ),
        )
        has_toolcalls, tc_list, assistant_msg = _extractToolCalls(response)
        @test has_toolcalls == true
        @test assistant_msg.api == "openai"
        @test assistant_msg.provider == "anthropic"
        @test assistant_msg.model == "claude-3-opus"
    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, assistant_msg = _extractToolCalls(response)
        @test has_toolcalls == true
        @test length(tc_list) == 1
        @test tc_list[1].name == "getWeather"
        @test tc_list[1].arguments["city"] == "Paris"
        # text block before tool_calls should be included in content
        @test length(assistant_msg.content) == 1
        @test assistant_msg.content[1].text == "Let me check."
    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, assistant_msg = _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

    @testset "content blocks with reasoning and text" begin
        response = Dict{String,Any}(
            "content" => Any[
                Dict{String,Any}("type" => "reasoning", "text" => "Thinking..."),
                Dict{String,Any}("type" => "text", "text" => "Here's the answer."),
            ],
        )
        has_toolcalls, tc_list, assistant_msg = _extractToolCalls(response)
        @test has_toolcalls == false
        @test length(tc_list) == 0
        @test length(assistant_msg.content) == 2
        @test assistant_msg.content[1].text == "Thinking..."
        @test assistant_msg.content[2].text == "Here's the answer."
    end

    @testset "content blocks with text and tool_call (tool_call not in content)" begin
        response = Dict{String,Any}(
            "content" => Any[
                Dict{String,Any}("type" => "text", "text" => "Sure, I'll check."),
                Dict{String,Any}(
                    "type" => "tool_call",
                    "id" => "tc_mix",
                    "name" => "getWeather",
                    "arguments" => Dict{String,Any}("city" => "London"),
                ),
            ],
        )
        has_toolcalls, tc_list, assistant_msg = _extractToolCalls(response)
        @test has_toolcalls == true
        @test length(tc_list) == 1
        @test tc_list[1].name == "getWeather"
        # text block included, tool_call block excluded from content
        @test length(assistant_msg.content) == 1
        @test assistant_msg.content[1].text == "Sure, I'll check."
    end

    # ------------------------------------------------------------------ #
    # assistantMessage construction                                       #
    # ------------------------------------------------------------------ #

    @testset "assistantMessage with error_message and errorMessage fallback" begin
        response = Dict{String,Any}(
            "error_message" => "rate limit",
            "content" => Any[Dict{String,Any}("type" => "text", "text" => "fail")],
        )
        has_toolcalls, tc_list, assistant_msg = _extractToolCalls(response)
        @test has_toolcalls == false
        @test assistant_msg.errorMessage == "rate limit"
    end

    @testset "assistantMessage with usage tracking" begin
        response = Dict{String,Any}(
            "content" => Any[Dict{String,Any}("type" => "text", "text" => "hi")],
            "usage" => llmUsage(100, 50),
        )
        has_toolcalls, tc_list, assistant_msg = _extractToolCalls(response)
        @test assistant_msg.usage.inputTokens == 100
        @test assistant_msg.usage.outputTokens == 50
    end

    @testset "assistantMessage with invalid usage defaults to zero" begin
        response = Dict{String,Any}(
            "content" => Any[Dict{String,Any}("type" => "text", "text" => "hi")],
            "usage" => "invalid",
        )
        has_toolcalls, tc_list, assistant_msg = _extractToolCalls(response)
        @test assistant_msg.usage.inputTokens == 0
        @test assistant_msg.usage.outputTokens == 0
    end

    @testset "reasoning_content as textContent" begin
        response = Dict{String,Any}(
            "reasoning_content" => textContent("internal thought"),
            "content" => Any[Dict{String,Any}("type" => "text", "text" => "output")],
        )
        has_toolcalls, tc_list, assistant_msg = _extractToolCalls(response)
        @test length(assistant_msg.content) == 2
        @test assistant_msg.content[1].text == "internal thought"
        @test assistant_msg.content[2].text == "output"
    end

    # ------------------------------------------------------------------ #
    # No tool call cases                                                  #
    # ------------------------------------------------------------------ #

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

    @testset "empty message" begin
        response = Dict{String,Any}()
        has_toolcalls, tc_list, assistant_msg = _extractToolCalls(response)
        @test has_toolcalls == false
        @test length(tc_list) == 0
        @test assistant_msg.role == "assistant"
        @test assistant_msg.stopReason == "end_turn"
        @test length(assistant_msg.content) == 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, assistant_msg = _extractToolCalls(response)
        @test has_toolcalls == false
        @test length(tc_list) == 0
        @test assistant_msg.role == "assistant"
    end

    @testset "message field is not a Dict" begin
        response = Dict{String,Any}(
            "message" => "not a dict",
            "content" => Any[Dict{String,Any}("type" => "text", "text" => "fallback")],
        )
        has_toolcalls, tc_list, assistant_msg = _extractToolCalls(response)
        @test has_toolcalls == false
        @test length(tc_list) == 0
        @test length(assistant_msg.content) == 1
    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, assistant_msg = _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, assistant_msg = _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, assistant_msg = _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, assistant_msg = _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, assistant_msg = _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

    @testset "tool_calls block with mixed content types (text + tool_calls)" begin
        response = Dict{String,Any}(
            "content" => Any[
                Dict{String,Any}("type" => "text", "text" => "I'll check both."),
                Dict{String,Any}(
                    "type" => "tool_calls",
                    "tool_calls" => Any[
                        Dict{String,Any}(
                            "type" => "function",
                            "function" => Dict{String,Any}(
                                "name" => "getWeather",
                                "arguments" => "{\"city\":\"London\"}",
                            ),
                            "id" => "tc_mix1",
                        ),
                        Dict{String,Any}(
                            "type" => "function",
                            "function" => Dict{String,Any}(
                                "name" => "getTime",
                                "arguments" => "{\"timezone\":\"UTC\"}",
                            ),
                            "id" => "tc_mix2",
                        ),
                    ],
                ),
            ],
        )
        has_toolcalls, tc_list, assistant_msg = _extractToolCalls(response)
        @test has_toolcalls == true
        @test length(tc_list) == 2
        @test tc_list[1].name == "getWeather"
        @test tc_list[2].name == "getTime"
        @test length(assistant_msg.content) == 1
        @test assistant_msg.content[1].text == "I'll check both."
    end

    @testset "tool_call block without arguments field" begin
        response = Dict{String,Any}(
            "content" => Any[
                Dict{String,Any}(
                    "type" => "tool_call",
                    "id" => "tc_noargs",
                    "name" => "getTime",
                ),
            ],
        )
        has_toolcalls, tc_list, assistant_msg = _extractToolCalls(response)
        @test has_toolcalls == true
        @test length(tc_list) == 1
        @test tc_list[1].arguments == Dict{String,Any}()
    end

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

    @testset "tool_calls block with non-AbstractDict elements" begin
        response = Dict{String,Any}(
            "content" => Any[
                Dict{String,Any}(
                    "type" => "tool_calls",
                    "tool_calls" => Any["not a dict", 42, nothing],
                ),
            ],
        )
        has_toolcalls, tc_list, assistant_msg = _extractToolCalls(response)
        @test has_toolcalls == false
        @test length(tc_list) == 0
    end

    @testset "content field is not a Vector" begin
        response = Dict{String,Any}(
            "content" => "not a vector",
        )
        has_toolcalls, tc_list, assistant_msg = _extractToolCalls(response)
        @test has_toolcalls == false
        @test length(tc_list) == 0
        @test length(assistant_msg.content) == 0
    end

    @testset "Dict-based response with all metadata fields" begin
        response = Dict{String,Any}(
            "api" => "openai",
            "provider" => "anthropic",
            "model" => "claude-3-sonnet",
            "content" => Any[
                Dict{String,Any}(
                    "type" => "tool_call",
                    "id" => "tc_meta",
                    "name" => "getTime",
                    "arguments" => Dict{String,Any}("city" => "Seoul"),
                ),
            ],
        )
        has_toolcalls, tc_list, assistant_msg = _extractToolCalls(response)
        @test has_toolcalls == true
        @test length(tc_list) == 1
        @test tc_list[1].name == "getTime"
        @test tc_list[1].arguments["city"] == "Seoul"
        @test assistant_msg.api == "openai"
        @test assistant_msg.provider == "anthropic"
        @test assistant_msg.model == "claude-3-sonnet"
    end

    @testset "default role is assistant" begin
        response = Dict{String,Any}(
            "content" => Any[Dict{String,Any}("type" => "text", "text" => "no role specified")],
        )
        has_toolcalls, tc_list, assistant_msg = _extractToolCalls(response)
        @test assistant_msg.role == "assistant"
    end

    @testset "tool call with custom role in message format" begin
        response = Dict{String,Any}(
            "message" => Dict{String,Any}(
                "role" => "custom_role",
                "tool_calls" => Any[
                    Dict{String,Any}(
                        "type" => "function",
                        "function" => Dict{String,Any}(
                            "name" => "getWeather",
                            "arguments" => "{}",
                        ),
                        "id" => "tc_role",
                    ),
                ],
            ),
        )
        has_toolcalls, tc_list, assistant_msg = _extractToolCalls(response)
        @test assistant_msg.role == "custom_role"
    end

    @testset "image content block handling" begin
        response = Dict{String,Any}(
            "content" => Any[
                Dict{String,Any}(
                    "type" => "image_url",
                    "image_url" => Dict("url" => "data:image/png;base64,abc123"),
                ),
                Dict{String,Any}("type" => "text", "text" => "What is this?"),
            ],
        )
        has_toolcalls, tc_list, assistant_msg = _extractToolCalls(response)
        @test has_toolcalls == false
        @test length(assistant_msg.content) == 2
        @test assistant_msg.content[1] isa textContent
        @test assistant_msg.content[1].text == ""
        @test assistant_msg.content[2].text == "What is this?"
    end

end