local PLUGIN_NAME = "ai-proxy"
local pl_file = require("pl.file")
local pl_replace = require("pl.stringx").replace
local cjson = require("cjson.safe")
local fmt = string.format
local llm = require("kong.llm")
local ai_shared = require("kong.llm.drivers.shared")

local SAMPLE_LLM_V2_CHAT_MULTIMODAL_IMAGE_URL = {
  messages = {
    {
      role = "user",
      content = {
        {
          type = "text",
          text = "What is in this picture?",
        },
        {
          type = "image_url",
          image_url = {
            url = "https://example.local/image.jpg",
          },
        },
      },
    },
    {
      role = "assistant",
      content = {
        {
          type = "text",
          text = "A picture of a cat.",
        },
      },
    },
    {
      role = "user",
      content = {
        {
          type = "text",
          text = "Now draw it wearing a party-hat.",
        },
      },
    },
  }
}

local SAMPLE_LLM_V2_CHAT_MULTIMODAL_IMAGE_B64 = {
  messages = {
    {
      role = "user",
      content = {
        {
          type = "text",
          text = "What is in this picture?",
        },
        {
          type = "image_url",
          image_url = {
            url = "data:image/png;base64,Y2F0X3BuZ19oZXJlX2xvbAo=",
          },
        },
      },
    },
    {
      role = "assistant",
      content = {
        {
          type = "text",
          text = "A picture of a cat.",
        },
      },
    },
    {
      role = "user",
      content = {
        {
          type = "text",
          text = "Now draw it wearing a party-hat.",
        },
      },
    },
  }
}

local SAMPLE_LLM_V1_CHAT = {
  messages = {
    [1] = {
      role = "system",
      content = "You are a mathematician."
    },
    [2] = {
      role = "assistant",
      content = "What is 1 + 1?"
    },
  },
}

local SAMPLE_LLM_V1_CHAT_WITH_SOME_OPTS = {
  messages = {
    [1] = {
      role = "system",
      content = "You are a mathematician."
    },
    [2] = {
      role = "assistant",
      content = "What is 1 + 1?"
    },
  },
  max_tokens = 256,
  temperature = 0.1,
  top_p = 0.2,
  some_extra_param = "string_val",
  another_extra_param = 0.5,
}

local SAMPLE_LLM_V1_CHAT_WITH_GUARDRAILS = {
  messages = {
    [1] = {
      role = "system",
      content = "You are a mathematician."
    },
    [2] = {
      role = "assistant",
      content = "What is 1 + 1?"
    },
  },
  guardrailConfig = {
    guardrailIdentifier = "yu5xwvfp4sud",
    guardrailVersion = "1",
    trace = "enabled",
  },
}

local SAMPLE_DOUBLE_FORMAT = {
  messages = {
    [1] = {
      role = "system",
      content = "You are a mathematician."
    },
    [2] = {
      role = "assistant",
      content = "What is 1 + 1?"
    },
  },
  prompt = "Hi world",
}

local SAMPLE_OPENAI_TOOLS_REQUEST = {
  messages = {
    [1] = {
      role = "user",
      content = "Is the NewPhone in stock?"
    },
  },
  tools = {
    [1] = {
      ['function'] = {
        parameters = {
          ['type'] = "object",
          properties = {
            product_name = {
              ['type'] = "string",
            },
          },
          required = {
            "product_name",
          },
        },
        name = "check_stock",
        description = "Check a product is in stock."
      },
      ['type'] = "function",
    },
  },
}

local SAMPLE_GEMINI_TOOLS_RESPONSE = {
  candidates = { {
    content = {
      role = "model",
      parts = { {
        functionCall = {
          name = "sql_execute",
          args = {
            product_name = "NewPhone"
          }
        }
      } }
    },
    finishReason = "STOP",
  } },
}

local SAMPLE_BEDROCK_TOOLS_RESPONSE = {
  metrics = {
    latencyMs = 3781
  },
  output = {
    message = {
      content = { {
        text = "Certainly! To calculate the sum of 121, 212, and 313, we can use the \"sumArea\" function that's available to us."
      }, {
        toolUse = {
          input = {
            areas = { 121, 212, 313 }
          },
          name = "sumArea",
          toolUseId = "tooluse_4ZakZPY9SiWoKWrAsY7_eg"
        }
      } },
      role = "assistant"
    }
  },
  stopReason = "tool_use",
  usage = {
    inputTokens = 410,
    outputTokens = 115,
    totalTokens = 525
  }
}

local FORMATS = {
  openai = {
    ["llm/v1/chat"] = {
      config = {
        name = "gpt-4",
        provider = "openai",
        options = {
          max_tokens = 512,
          temperature = 0.5,
        },
      },
    },
    ["llm/v1/completions"] = {
      config = {
        name = "gpt-3.5-turbo-instruct",
        provider = "openai",
        options = {
          max_tokens = 512,
          temperature = 0.5,
        },
      },
    },
  },
  cohere = {
    ["llm/v1/chat"] = {
      config = {
        name = "command",
        provider = "cohere",
        options = {
          max_tokens = 512,
          temperature = 0.5,
          top_p = 1.0
        },
      },
    },
    ["llm/v1/completions"] = {
      config = {
        name = "command",
        provider = "cohere",
        options = {
          max_tokens = 512,
          temperature = 0.5,
          top_p = 0.75,
          top_k = 5,
        },
      },
    },
  },
  anthropic = {
    ["llm/v1/chat"] = {
      config = {
        name = "claude-2.1",
        provider = "anthropic",
        options = {
          max_tokens = 512,
          temperature = 0.5,
          top_p = 1.0,
        },
      },
    },
    ["llm/v1/completions"] = {
      config = {
        name = "claude-2.1",
        provider = "anthropic",
        options = {
          max_tokens = 512,
          temperature = 0.5,
          top_p = 1.0,
        },
      },
    },
  },
  azure = {
    ["llm/v1/chat"] = {
      config = {
        name = "gpt-4",
        provider = "azure",
        options = {
          max_tokens = 512,
          temperature = 0.5,
          top_p = 1.0,
        },
      },
    },
    ["llm/v1/completions"] = {
      config = {
        name = "gpt-3.5-turbo-instruct",
      provider = "azure",
        options = {
          max_tokens = 512,
          temperature = 0.5,
          top_p = 1.0,
        },
      },
    },
  },
  llama2_raw = {
    ["llm/v1/chat"] = {
      config = {
        name = "llama2",
        provider = "llama2",
        options = {
          max_tokens = 512,
          temperature = 0.5,
          llama2_format = "raw",
          top_p = 1,
          top_k = 40,
        },
      },
    },
    ["llm/v1/completions"] = {
      config = {
        name = "llama2",
        provider = "llama2",
        options = {
          max_tokens = 512,
          temperature = 0.5,
          llama2_format = "raw",
        },
      },
    },
  },
  llama2_ollama = {
    ["llm/v1/chat"] = {
      config = {
        name = "llama2",
        provider = "llama2",
        options = {
          max_tokens = 512,
          temperature = 0.5,
          llama2_format = "ollama",
        },
      },
    },
    ["llm/v1/completions"] = {
      config = {
        name = "llama2",
        provider = "llama2",
        options = {
          max_tokens = 512,
          temperature = 0.5,
          llama2_format = "ollama",
        },
      },
    },
  },
  mistral_openai = {
    ["llm/v1/chat"] = {
      config = {
        name = "mistral-tiny",
        provider = "mistral",
        options = {
          max_tokens = 512,
          temperature = 0.5,
          mistral_format = "openai",
        },
      },
    },
  },
  mistral_ollama = {
    ["llm/v1/chat"] = {
      config = {
        name = "mistral-tiny",
        provider = "mistral",
        options = {
          max_tokens = 512,
          temperature = 0.5,
          mistral_format = "ollama",
        },
      },
    },
  },
  gemini = {
    ["llm/v1/chat"] = {
      config = {
        name = "gemini-pro",
        provider = "gemini",
        options = {
          max_tokens = 8192,
          temperature = 0.8,
          top_k = 1,
          top_p = 0.6,
        },
      },
    },
  },
  bedrock = {
    ["llm/v1/chat"] = {
      config = {
        name = "bedrock",
        provider = "bedrock",
        options = {
          max_tokens = 8192,
          temperature = 0.8,
          top_k = 1,
          top_p = 0.6,
        },
      },
    },
  },
}

local STREAMS = {
  openai = {
    ["llm/v1/chat"] = {
      name = "gpt-4",
      provider = "openai",
    },
    ["llm/v1/completions"] = {
      name = "gpt-3.5-turbo-instruct",
      provider = "openai",
    },
  },
  cohere = {
    ["llm/v1/chat"] = {
      name = "command",
      provider = "cohere",
    },
    ["llm/v1/completions"] = {
      name = "command-light",
      provider = "cohere",
    },
  },
}

local expected_stream_choices = {
  ["llm/v1/chat"] = {
    [1] = {
      delta = {
        content = "the answer",
      },
      finish_reason = ngx.null,
      index = 0,
      logprobs = ngx.null,
    },
  },
  ["llm/v1/completions"] = {
    [1] = {
      text = "the answer",
      finish_reason = ngx.null,
      index = 0,
      logprobs = ngx.null,
    },
  },
}

describe(PLUGIN_NAME .. ": (unit)", function()
  setup(function()
    package.loaded["kong.llm.drivers.shared"] = nil
    _G.TEST = true
    ai_shared = require("kong.llm.drivers.shared")
  end)

  teardown(function()
    _G.TEST = nil
  end)

  it("resolves referenceable plugin configuration from request context", function()
    local fake_request = {
      ["get_header"] = function(header_name)
        local headers = {
          ["from_header_1"] = "header_value_here_1",
          ["from_header_2"] = "header_value_here_2",
        }
        return headers[header_name]
      end,

      ["get_uri_captures"] = function()
        return {
          ["named"] = {
            ["uri_cap_1"] = "cap_value_here_1",
            ["uri_cap_2"] = "cap_value_here_2",
          },
        }
      end,

      ["get_query_arg"] = function(query_arg_name)
        local query_args = {
          ["arg_1"] = "arg_value_here_1",
          ["arg_2"] = "arg_value_here_2",
        }
        return query_args[query_arg_name]
      end,
    }

    local fake_config = {
      route_type = "llm/v1/chat",
      auth = {
        header_name = "api-key",
        header_value = "azure-key",
      },
      model = {
        name = "$(headers.from_header_1)",
        provider = "azure",
        options = {
          max_tokens = 256,
          temperature = 1.0,
          azure_instance = "$(uri_captures.uri_cap_1)",
          azure_deployment_id = "$(headers.from_header_1)",
          azure_api_version = "$(query_params.arg_1)",
          upstream_url = "https://$(uri_captures.uri_cap_1).example.com",
          bedrock = {
            aws_region = "$(uri_captures.uri_cap_1)",
          }
        },
      },
    }

    local result, err = ai_shared.merge_model_options(fake_request, fake_config)
    assert.is_falsy(err)
    assert.same(result.model.name, "header_value_here_1")
    assert.same(result.model.options, {
      azure_api_version = 'arg_value_here_1',
      azure_deployment_id = 'header_value_here_1',
      azure_instance = 'cap_value_here_1',
      max_tokens = 256,
      temperature = 1,
      upstream_url = "https://cap_value_here_1.example.com",
      bedrock = {
        aws_region = "cap_value_here_1",
      },
    })
  end)

  it("returns appropriate error when referenceable plugin configuration is missing from request context", function()
    local fake_request = {
      ["get_header"] = function(header_name)
        local headers = {
          ["from_header_1"] = "header_value_here_1",
          ["from_header_2"] = "header_value_here_2",
        }
        return headers[header_name]
      end,

      ["get_uri_captures"] = function()
        return {
          ["named"] = {
            ["uri_cap_1"] = "cap_value_here_1",
            ["uri_cap_2"] = "cap_value_here_2",
          },
        }
      end,

      ["get_query_arg"] = function(query_arg_name)
        local query_args = {
          ["arg_1"] = "arg_value_here_1",
          ["arg_2"] = "arg_value_here_2",
        }
        return query_args[query_arg_name]
      end,
    }

    local fake_config = {
      route_type = "llm/v1/chat",
      auth = {
        header_name = "api-key",
        header_value = "azure-key",
      },
      model = {
        name = "gpt-3.5-turbo",
        provider = "azure",
        options = {
          max_tokens = 256,
          temperature = 1.0,
          azure_instance = "$(uri_captures.uri_cap_3)",
          azure_deployment_id = "$(headers.from_header_1)",
          azure_api_version = "$(query_params.arg_1)",
        },
      },
    }

    local _, err = ai_shared.merge_model_options(fake_request, fake_config)
    assert.same("uri_captures key uri_cap_3 was not provided", err)

    local fake_config = {
      route_type = "llm/v1/chat",
      auth = {
        header_name = "api-key",
        header_value = "azure-key",
      },
      model = {
        name = "gpt-3.5-turbo",
        provider = "azure",
        options = {
          max_tokens = 256,
          temperature = 1.0,
          azure_instance = "$(uri_captures.uri_cap_1)",
          azure_deployment_id = "$(headers.from_header_1)",
          azure_api_version = "$(query_params.arg_1)",
          bedrock = {
            aws_region = "$(uri_captures.uri_cap_3)",
          }
        },
      },
    }

    local _, err = ai_shared.merge_model_options(fake_request, fake_config)
    assert.same("uri_captures key uri_cap_3 was not provided", err)

    local fake_config = {
      route_type = "llm/v1/chat",
      auth = {
        header_name = "api-key",
        header_value = "azure-key",
      },
      model = {
        name = "gpt-3.5-turbo",
        provider = "azure",
        options = {
          max_tokens = 256,
          temperature = 1.0,
          azure_instance = "$(uri_captures_uri_cap_1)",
        },
      },
    }

    local _, err = ai_shared.merge_model_options(fake_request, fake_config)
    assert.same("cannot parse expression for field '$(uri_captures_uri_cap_1)'", err)
  end)

  it("llm/v1/chat message is compatible with llm/v1/chat route", function()
    local compatible, err = llm.is_compatible(SAMPLE_LLM_V1_CHAT, "llm/v1/chat")

    assert.is_truthy(compatible)
    assert.is_nil(err)
  end)

  it("llm/v1/chat message is not compatible with llm/v1/completions route", function()
    local compatible, err = llm.is_compatible(SAMPLE_LLM_V1_CHAT, "llm/v1/completions")

    assert.is_falsy(compatible)
    assert.same("[llm/v1/chat] message format is not compatible with [llm/v1/completions] route type", err)
  end)

  it("double-format message is denied", function()
    local compatible, err = llm.is_compatible(SAMPLE_DOUBLE_FORMAT, "llm/v1/completions")

    assert.is_falsy(compatible)
    assert.same("request matches multiple LLM request formats", err)
  end)

  it("double-format message is denied", function()
    local compatible, err = llm.is_compatible(SAMPLE_DOUBLE_FORMAT, "llm/v1/completions")

    assert.is_falsy(compatible)
    assert.same("request matches multiple LLM request formats", err)
  end)

  for i, j in pairs(FORMATS) do

    describe(i .. " format tests", function()

      for k, l in pairs(j) do

        ---- actual testing code begins here
        describe(k .. " format test", function()

          local actual_request_table
          local driver = require("kong.llm.drivers." .. l.config.provider)


          -- what we do is first put the SAME request message from the user, through the converter, for this provider/format
          it("converts to provider request format correctly", function()
            -- load and check the driver
            assert(driver)

            -- load the standardised request, for this object type
            local request_json = pl_file.read(fmt("spec/fixtures/ai-proxy/unit/requests/%s.json", pl_replace(k, "/", "-")))
            local request_table, err = cjson.decode(request_json)
            assert.is_nil(err)

            -- send it
            local content_type, err
            actual_request_table, content_type, err = driver.to_format(request_table, l.config, k)
            assert.is_nil(err)
            assert.not_nil(content_type)

            -- load the expected outbound request to this provider
            local filename
            if l.config.provider == "llama2" then
              filename = fmt("spec/fixtures/ai-proxy/unit/expected-requests/%s/%s/%s.json", l.config.provider, l.config.options.llama2_format, pl_replace(k, "/", "-"))

            elseif l.config.provider == "mistral" then
              filename = fmt("spec/fixtures/ai-proxy/unit/expected-requests/%s/%s/%s.json", l.config.provider, l.config.options.mistral_format, pl_replace(k, "/", "-"))

            else
              filename = fmt("spec/fixtures/ai-proxy/unit/expected-requests/%s/%s.json", l.config.provider, pl_replace(k, "/", "-"))

            end

            local expected_request_json = pl_file.read(filename)
            local expected_request_table, err = cjson.decode(expected_request_json)
            assert.is_nil(err)

            -- compare the tables
            assert.same(expected_request_table, actual_request_table)
          end)


          -- then we put it through the converter that should come BACK from the provider, towards the user
          it("converts from provider response format correctly", function()
            -- load and check the driver
            assert(driver)

            -- load what the endpoint would really response with
            local filename
            if l.config.provider == "llama2" then
              filename = fmt("spec/fixtures/ai-proxy/unit/real-responses/%s/%s/%s.json", l.config.provider, l.config.options.llama2_format, pl_replace(k, "/", "-"))

            elseif l.config.provider == "mistral" then
              filename = fmt("spec/fixtures/ai-proxy/unit/real-responses/%s/%s/%s.json", l.config.provider, l.config.options.mistral_format, pl_replace(k, "/", "-"))

            else
              filename = fmt("spec/fixtures/ai-proxy/unit/real-responses/%s/%s.json", l.config.provider, pl_replace(k, "/", "-"))

            end
            local virtual_response_json = pl_file.read(filename)

            -- convert to kong format (emulate on response phase hook)
            local actual_response_json, err = driver.from_format(virtual_response_json, l.config, k)
            assert.is_nil(err)

            local actual_response_table, err = cjson.decode(actual_response_json)
            assert.is_nil(err)

            -- load the expected response body
            local filename
            if l.config.provider == "llama2" then
              filename = fmt("spec/fixtures/ai-proxy/unit/expected-responses/%s/%s/%s.json", l.config.provider, l.config.options.llama2_format, pl_replace(k, "/", "-"))

            elseif l.config.provider == "mistral" then
              filename = fmt("spec/fixtures/ai-proxy/unit/expected-responses/%s/%s/%s.json", l.config.provider, l.config.options.mistral_format, pl_replace(k, "/", "-"))

            else
              filename = fmt("spec/fixtures/ai-proxy/unit/expected-responses/%s/%s.json", l.config.provider, pl_replace(k, "/", "-"))

            end
            local expected_response_json = pl_file.read(filename)
            local expected_response_table, err = cjson.decode(expected_response_json)
            assert.is_nil(err)

            -- compare the tables
            assert.same(expected_response_table.choices[1].message, actual_response_table.choices[1].message)
            assert.same(actual_response_table.model, expected_response_table.model)
          end)
        end)
      end
    end)
  end

  -- streaming tests
  for provider_name, provider_format in pairs(STREAMS) do

    describe(provider_name .. " stream format tests", function()

      for format_name, config in pairs(provider_format) do

        ---- actual testing code begins here
        describe(format_name .. " format test", function()
          local driver = require("kong.llm.drivers." .. config.provider)

          -- what we do is first put the SAME request message from the user, through the converter, for this provider/format
          it("converts to provider request format correctly", function()
            -- load the real provider frame from file
            local real_stream_frame = pl_file.read(fmt("spec/fixtures/ai-proxy/unit/real-stream-frames/%s/%s.txt", config.provider, pl_replace(format_name, "/", "-")))

            -- use the shared function to produce an SSE format object
            local real_transformed_frame, err = ai_shared._frame_to_events(real_stream_frame)
            assert.is_nil(err)

            -- transform the SSE frame into OpenAI format
            real_transformed_frame, err = driver.from_format(real_transformed_frame[1], config, "stream/" .. format_name)
            assert.is_nil(err)
            real_transformed_frame, err = cjson.decode(real_transformed_frame)
            assert.is_nil(err)

            -- check it's what we expeced
            assert.same(expected_stream_choices[format_name], real_transformed_frame.choices)
          end)

        end)
      end
    end)

  end

  -- generic tests
  it("throws correct error when format is not supported", function()
    local driver = require("kong.llm.drivers.mistral")  -- one-shot, random example of provider with only prompt support

    local model_config = {
      route_type = "llm/v1/chatnopenotsupported",
      name = "mistral-tiny",
      provider = "mistral",
      options = {
        max_tokens = 512,
        temperature = 0.5,
        mistral_format = "ollama",
      },
    }

    local request_json = pl_file.read("spec/fixtures/ai-proxy/unit/requests/llm-v1-chat.json")
    local request_table, err = cjson.decode(request_json)
    assert.is_falsy(err)

    -- send it
    local actual_request_table, content_type, err = driver.to_format(request_table, model_config, model_config.route_type)
    assert.is_nil(actual_request_table)
    assert.is_nil(content_type)
    assert.equal(err, "no transformer available to format mistral://llm/v1/chatnopenotsupported/ollama")
  end)


  it("produces a correct default config merge", function()
    local formatted, err = ai_shared.merge_config_defaults(
      SAMPLE_LLM_V1_CHAT_WITH_SOME_OPTS,
      {
        max_tokens = 1024,
        top_p = 0.5,
      },
      "llm/v1/chat"
    )

    formatted.messages = nil  -- not needed for config merge

    assert.is_nil(err)
    assert.same({
      max_tokens          = 1024,
      temperature         = 0.1,
      top_p               = 0.5,
      some_extra_param    = "string_val",
      another_extra_param = 0.5,
    }, formatted)
  end)

  describe("streaming transformer tests", function()
    before_each(function()
      ai_shared._set_kong({
        ctx = {
          plugin = {}
        },
        log = {
          debug = function(...)
            print("[DEBUG] ", ...)
          end,
          err = function(...)
            print("[ERROR] ", ...)
          end,
        },
      })
    end)

    after_each(function()
      ai_shared._set_kong(nil)
    end)

    it("transforms Gemini type (split into two parts)", function()
      -- result
      local expected_result = pl_file.read(fmt("spec/fixtures/ai-proxy/unit/streaming-chunk-formats/partial-json-split/expected-output.bin"))

      -- body_filter 1
      local input = pl_file.read(fmt("spec/fixtures/ai-proxy/unit/streaming-chunk-formats/partial-json-split/input-1.bin"))
      local events_1 = ai_shared._frame_to_events(input, "application/json")

      -- body_filter 2
      input = pl_file.read(fmt("spec/fixtures/ai-proxy/unit/streaming-chunk-formats/partial-json-split/input-2.bin"))
      local events_2 = ai_shared._frame_to_events(input, "application/json")

      -- combine the two
      local result = ""
      for _, event_1 in ipairs(events_1) do
        result = result .. cjson.decode(event_1.data).candidates[1].content.parts[1].text
      end
      for _, event_2 in ipairs(events_2) do
        result = result .. cjson.decode(event_2.data).candidates[1].content.parts[1].text
      end

      assert.same(expected_result, result, true)
    end)

    it("transforms Gemini type (split into three parts)", function()
      -- result
      local expected_result = pl_file.read(fmt("spec/fixtures/ai-proxy/unit/streaming-chunk-formats/partial-json-split-three-parts/expected-output.bin"))

      -- body_filter 1
      local input = pl_file.read(fmt("spec/fixtures/ai-proxy/unit/streaming-chunk-formats/partial-json-split-three-parts/input-1.bin"))
      local events_1 = ai_shared._frame_to_events(input, "application/json")

      -- body_filter 2
      input = pl_file.read(fmt("spec/fixtures/ai-proxy/unit/streaming-chunk-formats/partial-json-split-three-parts/input-2.bin"))
      local events_2 = ai_shared._frame_to_events(input, "application/json")

      -- body_filter 3
      input = pl_file.read(fmt("spec/fixtures/ai-proxy/unit/streaming-chunk-formats/partial-json-split-three-parts/input-3.bin"))
      local events_3 = ai_shared._frame_to_events(input, "application/json")

      -- combine the two
      local result = ""
      for _, event_1 in ipairs(events_1) do
        result = result .. cjson.decode(event_1.data).candidates[1].content.parts[1].text
      end
      for _, event_2 in ipairs(events_2) do
        result = result .. cjson.decode(event_2.data).candidates[1].content.parts[1].text
      end
      for _, event_3 in ipairs(events_3) do
        result = result .. cjson.decode(event_3.data).candidates[1].content.parts[1].text
      end

      assert.same(expected_result, result, true)
    end)

    it("transforms Gemini type (beginning of stream)", function()
      local input = pl_file.read(fmt("spec/fixtures/ai-proxy/unit/streaming-chunk-formats/partial-json-beginning/input.bin"))
      local events = ai_shared._frame_to_events(input, "application/json")

      local expected = pl_file.read(fmt("spec/fixtures/ai-proxy/unit/streaming-chunk-formats/partial-json-beginning/expected-output.json"))
      local expected_events = cjson.decode(expected)

      assert.same(expected_events, events, true)
    end)

    it("transforms Gemini type (end of stream)", function()
      kong.ctx.plugin.gemini_state = {
        started = true,
        eof = false,
        input = "",
        pos = 1,
      }
      local input = pl_file.read(fmt("spec/fixtures/ai-proxy/unit/streaming-chunk-formats/partial-json-end/input.bin"))
      local events = ai_shared._frame_to_events(input, "application/json")

      local expected = pl_file.read(fmt("spec/fixtures/ai-proxy/unit/streaming-chunk-formats/partial-json-end/expected-output.json"))
      local expected_events = cjson.decode(expected)

      assert.same(expected_events, events, true)
    end)

    it("transforms complete-json type", function()
      local input = pl_file.read(fmt("spec/fixtures/ai-proxy/unit/streaming-chunk-formats/complete-json/input.bin"))
      local events = ai_shared._frame_to_events(input, "text/event-stream")  -- not "truncated json mode" like Gemini

      local expected = pl_file.read(fmt("spec/fixtures/ai-proxy/unit/streaming-chunk-formats/complete-json/expected-output.json"))
      local expected_events = cjson.decode(expected)

      assert.same(expected_events, events)
    end)

    it("transforms text/event-stream type", function()
      local input = pl_file.read(fmt("spec/fixtures/ai-proxy/unit/streaming-chunk-formats/text-event-stream/input.bin"))
      local events = ai_shared._frame_to_events(input, "text/event-stream")  -- not "truncated json mode" like Gemini

      local expected = pl_file.read(fmt("spec/fixtures/ai-proxy/unit/streaming-chunk-formats/text-event-stream/expected-output.json"))
      local expected_events = cjson.decode(expected)

      assert.same(expected_events, events)
    end)

    it("transforms application/vnd.amazon.eventstream (AWS) type", function()
      local input = pl_file.read(fmt("spec/fixtures/ai-proxy/unit/streaming-chunk-formats/aws/input.bin"))
      local events = ai_shared._frame_to_events(input, "application/vnd.amazon.eventstream")

      local expected = pl_file.read(fmt("spec/fixtures/ai-proxy/unit/streaming-chunk-formats/aws/expected-output.json"))
      local expected_events = cjson.decode(expected)

      assert.equal(#events, #expected_events)
      for i, _ in ipairs(expected_events) do
        -- tables are random ordered, so we need to compare each serialized event
        assert.same(cjson.decode(events[i].data), cjson.decode(expected_events[i].data))
      end
    end)

  end)

  describe("count_words", function()
    local c = ai_shared._count_words

    it("normal prompts", function()
      assert.same(10, c(string.rep("apple ", 10)))
    end)

    it("multi-modal prompts", function()
      assert.same(10, c({
        {
          type = "text",
          text = string.rep("apple ", 10),
        },
      }))

      assert.same(20, c({
        {
          type = "text",
          text = string.rep("apple ", 10),
        },
        {
          type = "text",
          text = string.rep("banana ", 10),
        },
      }))

      assert.same(10, c({
        {
          type = "not_text",
          text = string.rep("apple ", 10),
        },
        {
          type = "text",
          text = string.rep("banana ", 10),
        },
        {
          type = "text",
          -- somehow malformed
        },
      }))
    end)
  end)

  describe("gemini multimodal", function()
    local gemini_driver

    setup(function()
      _G._TEST = true
      package.loaded["kong.llm.drivers.gemini"] = nil
      gemini_driver = require("kong.llm.drivers.gemini")
    end)

    teardown(function()
      _G._TEST = nil
    end)

    it("transforms a text type prompt to gemini GOOD", function()
      local gemini_prompt, err = gemini_driver._openai_part_to_gemini_part(
        {
          ["type"] = "text",
          ["text"] = "What is in this picture?",
        })

      assert.not_nil(gemini_prompt)
      assert.is_nil(err)

      assert.same(gemini_prompt,
        {
          ["text"] = "What is in this picture?",
        })
    end)

    it("transforms a text type prompt to gemini BAD MISSING TEXT FIELD", function()
      local gemini_prompt, err = gemini_driver._openai_part_to_gemini_part(
        {
          ["type"] = "text",
          ["bad_text_field"] = "What is in this picture?",
        })

      assert.is_nil(gemini_prompt)
      assert.not_nil(err)

      assert.same("message part type is 'text' but is missing .text block", err)
    end)

    it("transforms an image_url type prompt when data is a URL to gemini GOOD", function()
      local gemini_prompt, err = gemini_driver._openai_part_to_gemini_part(
        {
          ["type"] = "image_url",
          ["image_url"] = {
            ["url"] = "https://example.local/image.jpg",
          },
        })

      assert.not_nil(gemini_prompt)
      assert.is_nil(err)

      assert.same(gemini_prompt,
        {
          ["fileData"] = {
            ["fileUri"] = "https://example.local/image.jpg",
            ["mimeType"] = "image/generic",
          },
        })
    end)

    it("transforms an image_url type prompt when data is a URL to gemini BAD MISSING IMAGE FIELD", function()
      local gemini_prompt, err = gemini_driver._openai_part_to_gemini_part(
        {
          ["type"] = "image_url",
          ["image_url"] = "https://example.local/image.jpg",
        })

      assert.is_nil(gemini_prompt)
      assert.not_nil(err)

      assert.same("message part type is 'image_url' but is missing .image_url.url block", err)
    end)

    it("fails to transform a non-mapped multimodal entity type", function()
      local gemini_prompt, err = gemini_driver._openai_part_to_gemini_part(
        {
          ["type"] = "doesnt_exist",
          ["doesnt_exist"] = "https://example.local/video.mp4",
        })

      assert.is_nil(gemini_prompt)
      assert.not_nil(err)

      assert.same("cannot transform part of type 'doesnt_exist' to Gemini format", err)
    end)

    it("transforms 'describe this image' via URL from openai to gemini", function()
      local gemini_prompt, _, err = gemini_driver._to_gemini_chat_openai(SAMPLE_LLM_V2_CHAT_MULTIMODAL_IMAGE_URL)

      assert.is_nil(err)
      assert.not_nil(gemini_prompt)

      gemini_prompt.generationConfig = nil  -- not needed for comparison

      assert.same({
        ["contents"] = {
          {
            ["role"] = "user",
            ["parts"] = {
              {
                ["text"] = "What is in this picture?",
              },
              {
                ["fileData"] = {
                  ["fileUri"] = "https://example.local/image.jpg",
                  ["mimeType"] = "image/generic",
                },
              }
            },
          },
          {
            ["role"] = "model",
            ["parts"] = {
              {
                ["text"] = "A picture of a cat.",
              },
            },
          },
          {
            ["role"] = "user",
            ["parts"] = {
              {
                ["text"] = "Now draw it wearing a party-hat.",
              },
            },
          },
        }
      }, gemini_prompt)
    end)

    it("transforms 'describe this image' via base64 from openai to gemini", function()
      local gemini_prompt, _, err = gemini_driver._to_gemini_chat_openai(SAMPLE_LLM_V2_CHAT_MULTIMODAL_IMAGE_B64)

      assert.is_nil(err)
      assert.not_nil(gemini_prompt)

      gemini_prompt.generationConfig = nil  -- not needed for comparison

      assert.same({
        ["contents"] = {
          {
            ["role"] = "user",
            ["parts"] = {
              {
                ["text"] = "What is in this picture?",
              },
              {
                ["inlineData"] = {
                  ["data"] = "Y2F0X3BuZ19oZXJlX2xvbAo=",
                  ["mimeType"] = "image/png",
                },
              }
            },
          },
          {
            ["role"] = "model",
            ["parts"] = {
              {
                ["text"] = "A picture of a cat.",
              },
            },
          },
          {
            ["role"] = "user",
            ["parts"] = {
              {
                ["text"] = "Now draw it wearing a party-hat.",
              },
            },
          },
        }
      }, gemini_prompt)
    end)

  end)


  describe("gemini tools", function()
    local gemini_driver

    setup(function()
      _G._TEST = true
      package.loaded["kong.llm.drivers.gemini"] = nil
      gemini_driver = require("kong.llm.drivers.gemini")
    end)

    teardown(function()
      _G._TEST = nil
    end)

    it("transforms openai tools to gemini tools GOOD", function()
      local gemini_tools = gemini_driver._to_tools(SAMPLE_OPENAI_TOOLS_REQUEST.tools)

      assert.not_nil(gemini_tools)
      assert.same(gemini_tools, {
        {
          function_declarations = {
            {
              description = "Check a product is in stock.",
              name = "check_stock",
              parameters = {
                properties = {
                  product_name = {
                    type = "string"
                  }
                },
                required = {
                  "product_name"
                },
                type = "object"
              }
            }
          }
        }
      })
    end)

    it("transforms openai tools to gemini tools NO_TOOLS", function()
      local gemini_tools = gemini_driver._to_tools(SAMPLE_LLM_V1_CHAT)

      assert.is_nil(gemini_tools)
    end)

    it("transforms openai tools to gemini tools NIL", function()
      local gemini_tools = gemini_driver._to_tools(nil)

      assert.is_nil(gemini_tools)
    end)

    it("transforms gemini tools to openai tools GOOD", function()
      local openai_tools = gemini_driver._from_gemini_chat_openai(SAMPLE_GEMINI_TOOLS_RESPONSE, {}, "llm/v1/chat")

      assert.not_nil(openai_tools)

      openai_tools = cjson.decode(openai_tools)
      assert.same(openai_tools.choices[1].message.tool_calls[1]['function'], {
        name = "sql_execute",
        arguments = "{\"product_name\":\"NewPhone\"}"
      })
    end)
  end)

  describe("bedrock tools", function()
    local bedrock_driver

    setup(function()
      _G._TEST = true
      package.loaded["kong.llm.drivers.bedrock"] = nil
      bedrock_driver = require("kong.llm.drivers.bedrock")
    end)

    teardown(function()
      _G._TEST = nil
    end)

    it("transforms openai tools to bedrock tools GOOD", function()
      local bedrock_tools = bedrock_driver._to_tools(SAMPLE_OPENAI_TOOLS_REQUEST.tools)

      assert.not_nil(bedrock_tools)
      assert.same(bedrock_tools, {
        {
          toolSpec = {
            description = "Check a product is in stock.",
            inputSchema = {
              json = {
                properties = {
                  product_name = {
                    type = "string"
                  }
                },
                required = {
                  "product_name"
                },
                type = "object"
              }
            },
            name = "check_stock"
          }
        }
      })
    end)

    it("transforms openai tools to bedrock tools NO_TOOLS", function()
      local bedrock_tools = bedrock_driver._to_tools(SAMPLE_LLM_V1_CHAT)

      assert.is_nil(bedrock_tools)
    end)

    it("transforms openai tools to bedrock tools NIL", function()
      local bedrock_tools = bedrock_driver._to_tools(nil)

      assert.is_nil(bedrock_tools)
    end)

    it("transforms bedrock tools to openai tools GOOD", function()
      local openai_tools = bedrock_driver._from_tool_call_response(SAMPLE_BEDROCK_TOOLS_RESPONSE.output.message.content)

      assert.not_nil(openai_tools)

      assert.same(openai_tools[1]['function'], {
        name = "sumArea",
        arguments = "{\"areas\":[121,212,313]}"
      })
    end)

    it("transforms guardrails into bedrock generation config", function()
      local model_info = {
        route_type = "llm/v1/chat",
        name = "some-model",
        provider = "bedrock",
      }
      local bedrock_guardrails = bedrock_driver._to_bedrock_chat_openai(SAMPLE_LLM_V1_CHAT_WITH_GUARDRAILS, model_info, "llm/v1/chat")

      assert.not_nil(bedrock_guardrails)

      assert.same(bedrock_guardrails.guardrailConfig, {
        ['guardrailIdentifier'] = 'yu5xwvfp4sud',
        ['guardrailVersion'] = '1',
        ['trace'] = 'enabled',
      })
    end)
  end)
end)


describe(PLUGIN_NAME .. ": (unit)", function()
  setup(function()
    package.loaded["kong.llm.drivers.shared"] = nil
    _G.TEST = true
    ai_shared = require("kong.llm.drivers.shared")
  end)

  teardown(function()
    _G.TEST = nil
  end)

  it("resolves referenceable plugin configuration from request context", function()
    local fake_request = {
      ["get_header"] = function(header_name)
        local headers = {
          ["from_header_1"] = "header_value_here_1",
          ["from_header_2"] = "header_value_here_2",
        }
        return headers[header_name]
      end,

      ["get_uri_captures"] = function()
        return {
          ["named"] = {
            ["uri_cap_1"] = "cap_value_here_1",
            ["uri_cap_2"] = "cap_value_here_2",
          },
        }
      end,

      ["get_query_arg"] = function(query_arg_name)
        local query_args = {
          ["arg_1"] = "arg_value_here_1",
          ["arg_2"] = "arg_value_here_2",
        }
        return query_args[query_arg_name]
      end,
    }

    local fake_config = {
      route_type = "llm/v1/chat",
      auth = {
        header_name = "api-key",
        header_value = "azure-key",
      },
      model = {
        name = "gpt-3.5-turbo",
        provider = "azure",
        options = {
          max_tokens = 256,
          temperature = 1.0,
          azure_instance = "$(uri_captures.uri_cap_1)",
          azure_deployment_id = "$(headers.from_header_1)",
          azure_api_version = "$(query_params.arg_1)",
        },
      },
    }

    local result, err = ai_shared.merge_model_options(fake_request, fake_config)
    assert.is_falsy(err)
    assert.same(result.model.options, {
      ['azure_api_version'] = 'arg_value_here_1',
      ['azure_deployment_id'] = 'header_value_here_1',
      ['azure_instance'] = 'cap_value_here_1',
      ['max_tokens'] = 256,
      ['temperature'] = 1,
    })
  end)

  it("resolves referenceable model name from request context", function()
    local fake_request = {
      ["get_header"] = function(header_name)
        local headers = {
          ["from_header_1"] = "header_value_here_1",
          ["from_header_2"] = "header_value_here_2",
        }
        return headers[header_name]
      end,

      ["get_uri_captures"] = function()
        return {
          ["named"] = {
            ["uri_cap_1"] = "cap_value_here_1",
            ["uri_cap_2"] = "cap_value_here_2",
          },
        }
      end,

      ["get_query_arg"] = function(query_arg_name)
        local query_args = {
          ["arg_1"] = "arg_value_here_1",
          ["arg_2"] = "arg_value_here_2",
        }
        return query_args[query_arg_name]
      end,
    }

    local fake_config = {
      route_type = "llm/v1/chat",
      auth = {
        header_name = "api-key",
        header_value = "azure-key",
      },
      model = {
        name = "$(uri_captures.uri_cap_2)",
        provider = "azure",
        options = {
          max_tokens = 256,
          temperature = 1.0,
          azure_instance = "string-1",
          azure_deployment_id = "string-2",
          azure_api_version = "string-3",
        },
      },
    }

    local result, err = ai_shared.merge_model_options(fake_request, fake_config)
    assert.is_falsy(err)
    assert.same("cap_value_here_2", result.model.name)
  end)

  it("returns appropriate error when referenceable plugin configuration is missing from request context", function()
    local fake_request = {
      ["get_header"] = function(header_name)
        local headers = {
          ["from_header_1"] = "header_value_here_1",
          ["from_header_2"] = "header_value_here_2",
        }
        return headers[header_name]
      end,

      ["get_uri_captures"] = function()
        return {
          ["named"] = {
            ["uri_cap_1"] = "cap_value_here_1",
            ["uri_cap_2"] = "cap_value_here_2",
          },
        }
      end,

      ["get_query_arg"] = function(query_arg_name)
        local query_args = {
          ["arg_1"] = "arg_value_here_1",
          ["arg_2"] = "arg_value_here_2",
        }
        return query_args[query_arg_name]
      end,
    }

    local fake_config = {
      route_type = "llm/v1/chat",
      auth = {
        header_name = "api-key",
        header_value = "azure-key",
      },
      model = {
        name = "gpt-3.5-turbo",
        provider = "azure",
        options = {
          max_tokens = 256,
          temperature = 1.0,
          azure_instance = "$(uri_captures.uri_cap_3)",
          azure_deployment_id = "$(headers.from_header_1)",
          azure_api_version = "$(query_params.arg_1)",
        },
      },
    }

    local _, err = ai_shared.merge_model_options(fake_request, fake_config)
    assert.same("uri_captures key uri_cap_3 was not provided", err)
  end)

  it("llm/v1/chat message is compatible with llm/v1/chat route", function()
    local compatible, err = llm.is_compatible(SAMPLE_LLM_V1_CHAT, "llm/v1/chat")

    assert.is_truthy(compatible)
    assert.is_nil(err)
  end)

  it("llm/v1/chat message is not compatible with llm/v1/completions route", function()
    local compatible, err = llm.is_compatible(SAMPLE_LLM_V1_CHAT, "llm/v1/completions")

    assert.is_falsy(compatible)
    assert.same("[llm/v1/chat] message format is not compatible with [llm/v1/completions] route type", err)
  end)

  it("double-format message is denied", function()
    local compatible, err = llm.is_compatible(SAMPLE_DOUBLE_FORMAT, "llm/v1/completions")

    assert.is_falsy(compatible)
    assert.same("request matches multiple LLM request formats", err)
  end)

  it("double-format message is denied", function()
    local compatible, err = llm.is_compatible(SAMPLE_DOUBLE_FORMAT, "llm/v1/completions")

    assert.is_falsy(compatible)
    assert.same("request matches multiple LLM request formats", err)
  end)

  for i, j in pairs(FORMATS) do

    describe(i .. " format tests", function()

      for k, l in pairs(j) do

        ---- actual testing code begins here
        describe(k .. " format test", function()

          local actual_request_table
          local driver = require("kong.llm.drivers." .. l.config.provider)


          -- what we do is first put the SAME request message from the user, through the converter, for this provider/format
          it("converts to provider request format correctly", function()
            -- load and check the driver
            assert(driver)

            -- load the standardised request, for this object type
            local request_json = pl_file.read(fmt("spec/fixtures/ai-proxy/unit/requests/%s.json", pl_replace(k, "/", "-")))
            local request_table, err = cjson.decode(request_json)
            assert.is_nil(err)

            -- send it
            local content_type, err
            actual_request_table, content_type, err = driver.to_format(request_table, l.config, k)
            assert.is_nil(err)
            assert.not_nil(content_type)

            -- load the expected outbound request to this provider
            local filename
            if l.config.provider == "llama2" then
              filename = fmt("spec/fixtures/ai-proxy/unit/expected-requests/%s/%s/%s.json", l.config.provider, l.config.options.llama2_format, pl_replace(k, "/", "-"))

            elseif l.config.provider == "mistral" then
              filename = fmt("spec/fixtures/ai-proxy/unit/expected-requests/%s/%s/%s.json", l.config.provider, l.config.options.mistral_format, pl_replace(k, "/", "-"))

            else
              filename = fmt("spec/fixtures/ai-proxy/unit/expected-requests/%s/%s.json", l.config.provider, pl_replace(k, "/", "-"))

            end

            local expected_request_json = pl_file.read(filename)
            local expected_request_table, err = cjson.decode(expected_request_json)
            assert.is_nil(err)

            -- compare the tables
            assert.same(expected_request_table, actual_request_table)
          end)


          -- then we put it through the converter that should come BACK from the provider, towards the user
          it("converts from provider response format correctly", function()
            -- load and check the driver
            assert(driver)

            -- load what the endpoint would really response with
            local filename
            if l.config.provider == "llama2" then
              filename = fmt("spec/fixtures/ai-proxy/unit/real-responses/%s/%s/%s.json", l.config.provider, l.config.options.llama2_format, pl_replace(k, "/", "-"))

            elseif l.config.provider == "mistral" then
              filename = fmt("spec/fixtures/ai-proxy/unit/real-responses/%s/%s/%s.json", l.config.provider, l.config.options.mistral_format, pl_replace(k, "/", "-"))

            else
              filename = fmt("spec/fixtures/ai-proxy/unit/real-responses/%s/%s.json", l.config.provider, pl_replace(k, "/", "-"))

            end
            local virtual_response_json = pl_file.read(filename)

            -- convert to kong format (emulate on response phase hook)
            local actual_response_json, err = driver.from_format(virtual_response_json, l.config, k)
            assert.is_nil(err)

            local actual_response_table, err = cjson.decode(actual_response_json)
            assert.is_nil(err)

            -- load the expected response body
            local filename
            if l.config.provider == "llama2" then
              filename = fmt("spec/fixtures/ai-proxy/unit/expected-responses/%s/%s/%s.json", l.config.provider, l.config.options.llama2_format, pl_replace(k, "/", "-"))

            elseif l.config.provider == "mistral" then
              filename = fmt("spec/fixtures/ai-proxy/unit/expected-responses/%s/%s/%s.json", l.config.provider, l.config.options.mistral_format, pl_replace(k, "/", "-"))

            else
              filename = fmt("spec/fixtures/ai-proxy/unit/expected-responses/%s/%s.json", l.config.provider, pl_replace(k, "/", "-"))

            end
            local expected_response_json = pl_file.read(filename)
            local expected_response_table, err = cjson.decode(expected_response_json)
            assert.is_nil(err)

            -- compare the tables
            assert.same(expected_response_table.choices[1].message, actual_response_table.choices[1].message)
            assert.same(actual_response_table.model, expected_response_table.model)
          end)
        end)
      end
    end)
  end

  -- streaming tests
  for provider_name, provider_format in pairs(STREAMS) do

    describe(provider_name .. " stream format tests", function()

      for format_name, config in pairs(provider_format) do

        ---- actual testing code begins here
        describe(format_name .. " format test", function()
          local driver = require("kong.llm.drivers." .. config.provider)

          -- what we do is first put the SAME request message from the user, through the converter, for this provider/format
          it("converts to provider request format correctly", function()
            -- load the real provider frame from file
            local real_stream_frame = pl_file.read(fmt("spec/fixtures/ai-proxy/unit/real-stream-frames/%s/%s.txt", config.provider, pl_replace(format_name, "/", "-")))

            -- use the shared function to produce an SSE format object
            local real_transformed_frame, err = ai_shared._frame_to_events(real_stream_frame)
            assert.is_nil(err)

            -- transform the SSE frame into OpenAI format
            real_transformed_frame, err = driver.from_format(real_transformed_frame[1], config, "stream/" .. format_name)
            assert.is_nil(err)
            real_transformed_frame, err = cjson.decode(real_transformed_frame)
            assert.is_nil(err)

            -- check it's what we expeced
            assert.same(expected_stream_choices[format_name], real_transformed_frame.choices)
          end)

        end)
      end
    end)

  end

  -- generic tests
  it("throws correct error when format is not supported", function()
    local driver = require("kong.llm.drivers.mistral")  -- one-shot, random example of provider with only prompt support

    local model_config = {
      route_type = "llm/v1/chatnopenotsupported",
      name = "mistral-tiny",
      provider = "mistral",
      options = {
        max_tokens = 512,
        temperature = 0.5,
        mistral_format = "ollama",
      },
    }

    local request_json = pl_file.read("spec/fixtures/ai-proxy/unit/requests/llm-v1-chat.json")
    local request_table, err = cjson.decode(request_json)
    assert.is_falsy(err)

    -- send it
    local actual_request_table, content_type, err = driver.to_format(request_table, model_config, model_config.route_type)
    assert.is_nil(actual_request_table)
    assert.is_nil(content_type)
    assert.equal(err, "no transformer available to format mistral://llm/v1/chatnopenotsupported/ollama")
  end)


  it("produces a correct default config merge", function()
    local formatted, err = ai_shared.merge_config_defaults(
      SAMPLE_LLM_V1_CHAT_WITH_SOME_OPTS,
      {
        max_tokens = 1024,
        top_p = 0.5,
      },
      "llm/v1/chat"
    )

    formatted.messages = nil  -- not needed for config merge

    assert.is_nil(err)
    assert.same({
      max_tokens          = 1024,
      temperature         = 0.1,
      top_p               = 0.5,
      some_extra_param    = "string_val",
      another_extra_param = 0.5,
    }, formatted)
  end)

  describe("streaming transformer tests", function()
    before_each(function()
      ai_shared._set_kong({
        ctx = {
          plugin = {}
        },
        log = {
          debug = function(...)
            print("[DEBUG] ", ...)
          end,
          err = function(...)
            print("[ERROR] ", ...)
          end,
        },
      })
    end)

    after_each(function()
      ai_shared._set_kong(nil)
    end)

    it("transforms Gemini type (split into two parts)", function()
      -- result
      local expected_result = pl_file.read(fmt("spec/fixtures/ai-proxy/unit/streaming-chunk-formats/partial-json-split/expected-output.bin"))

      -- body_filter 1
      local input = pl_file.read(fmt("spec/fixtures/ai-proxy/unit/streaming-chunk-formats/partial-json-split/input-1.bin"))
      local events_1 = ai_shared._frame_to_events(input, "application/json")

      -- body_filter 2
      input = pl_file.read(fmt("spec/fixtures/ai-proxy/unit/streaming-chunk-formats/partial-json-split/input-2.bin"))
      local events_2 = ai_shared._frame_to_events(input, "application/json")

      -- combine the two
      local result = ""
      for _, event_1 in ipairs(events_1) do
        result = result .. cjson.decode(event_1.data).candidates[1].content.parts[1].text
      end
      for _, event_2 in ipairs(events_2) do
        result = result .. cjson.decode(event_2.data).candidates[1].content.parts[1].text
      end

      assert.same(expected_result, result, true)
    end)

    it("transforms Gemini type (split into three parts)", function()
      -- result
      local expected_result = pl_file.read(fmt("spec/fixtures/ai-proxy/unit/streaming-chunk-formats/partial-json-split-three-parts/expected-output.bin"))

      -- body_filter 1
      local input = pl_file.read(fmt("spec/fixtures/ai-proxy/unit/streaming-chunk-formats/partial-json-split-three-parts/input-1.bin"))
      local events_1 = ai_shared._frame_to_events(input, "application/json")

      -- body_filter 2
      input = pl_file.read(fmt("spec/fixtures/ai-proxy/unit/streaming-chunk-formats/partial-json-split-three-parts/input-2.bin"))
      local events_2 = ai_shared._frame_to_events(input, "application/json")

      -- body_filter 3
      input = pl_file.read(fmt("spec/fixtures/ai-proxy/unit/streaming-chunk-formats/partial-json-split-three-parts/input-3.bin"))
      local events_3 = ai_shared._frame_to_events(input, "application/json")

      -- combine the two
      local result = ""
      for _, event_1 in ipairs(events_1) do
        result = result .. cjson.decode(event_1.data).candidates[1].content.parts[1].text
      end
      for _, event_2 in ipairs(events_2) do
        result = result .. cjson.decode(event_2.data).candidates[1].content.parts[1].text
      end
      for _, event_3 in ipairs(events_3) do
        result = result .. cjson.decode(event_3.data).candidates[1].content.parts[1].text
      end

      assert.same(expected_result, result, true)
    end)

    it("transforms Gemini type (beginning of stream)", function()
      local input = pl_file.read(fmt("spec/fixtures/ai-proxy/unit/streaming-chunk-formats/partial-json-beginning/input.bin"))
      local events = ai_shared._frame_to_events(input, "application/json")

      local expected = pl_file.read(fmt("spec/fixtures/ai-proxy/unit/streaming-chunk-formats/partial-json-beginning/expected-output.json"))
      local expected_events = cjson.decode(expected)

      assert.same(expected_events, events, true)
    end)

    it("transforms Gemini type (end of stream)", function()
      kong.ctx.plugin.gemini_state = {
        started = true,
        eof = false,
        input = "",
        pos = 1,
      }
      local input = pl_file.read(fmt("spec/fixtures/ai-proxy/unit/streaming-chunk-formats/partial-json-end/input.bin"))
      local events = ai_shared._frame_to_events(input, "application/json")

      local expected = pl_file.read(fmt("spec/fixtures/ai-proxy/unit/streaming-chunk-formats/partial-json-end/expected-output.json"))
      local expected_events = cjson.decode(expected)

      assert.same(expected_events, events, true)
    end)

    it("transforms complete-json type", function()
      local input = pl_file.read(fmt("spec/fixtures/ai-proxy/unit/streaming-chunk-formats/complete-json/input.bin"))
      local events = ai_shared._frame_to_events(input, "text/event-stream")  -- not "truncated json mode" like Gemini

      local expected = pl_file.read(fmt("spec/fixtures/ai-proxy/unit/streaming-chunk-formats/complete-json/expected-output.json"))
      local expected_events = cjson.decode(expected)

      assert.same(expected_events, events)
    end)

    it("transforms text/event-stream type", function()
      local input = pl_file.read(fmt("spec/fixtures/ai-proxy/unit/streaming-chunk-formats/text-event-stream/input.bin"))
      local events = ai_shared._frame_to_events(input, "text/event-stream")  -- not "truncated json mode" like Gemini

      local expected = pl_file.read(fmt("spec/fixtures/ai-proxy/unit/streaming-chunk-formats/text-event-stream/expected-output.json"))
      local expected_events = cjson.decode(expected)

      assert.same(expected_events, events)
    end)

    it("transforms application/vnd.amazon.eventstream (AWS) type", function()
      local input = pl_file.read(fmt("spec/fixtures/ai-proxy/unit/streaming-chunk-formats/aws/input.bin"))
      local events = ai_shared._frame_to_events(input, "application/vnd.amazon.eventstream")

      local expected = pl_file.read(fmt("spec/fixtures/ai-proxy/unit/streaming-chunk-formats/aws/expected-output.json"))
      local expected_events = cjson.decode(expected)

      assert.equal(#events, #expected_events)
      for i, _ in ipairs(expected_events) do
        -- tables are random ordered, so we need to compare each serialized event
        assert.same(cjson.decode(events[i].data), cjson.decode(expected_events[i].data))
      end
    end)

  end)

  describe("count_words", function()
    local c = ai_shared._count_words

    it("normal prompts", function()
      assert.same(10, c(string.rep("apple ", 10)))
    end)

    it("multi-modal prompts", function()
      assert.same(10, c({
        {
          type = "text",
          text = string.rep("apple ", 10),
        },
      }))

      assert.same(20, c({
        {
          type = "text",
          text = string.rep("apple ", 10),
        },
        {
          type = "text",
          text = string.rep("banana ", 10),
        },
      }))

      assert.same(10, c({
        {
          type = "not_text",
          text = string.rep("apple ", 10),
        },
        {
          type = "text",
          text = string.rep("banana ", 10),
        },
        {
          type = "text",
          -- somehow malformed
        },
      }))
    end)
  end)

  describe("gemini multimodal", function()
    local gemini_driver

    setup(function()
      _G._TEST = true
      package.loaded["kong.llm.drivers.gemini"] = nil
      gemini_driver = require("kong.llm.drivers.gemini")
    end)

    teardown(function()
      _G._TEST = nil
    end)

    it("transforms a text type prompt to gemini GOOD", function()
      local gemini_prompt, err = gemini_driver._openai_part_to_gemini_part(
        {
          ["type"] = "text",
          ["text"] = "What is in this picture?",
        })

      assert.not_nil(gemini_prompt)
      assert.is_nil(err)

      assert.same(gemini_prompt,
        {
          ["text"] = "What is in this picture?",
        })
    end)

    it("transforms a text type prompt to gemini BAD MISSING TEXT FIELD", function()
      local gemini_prompt, err = gemini_driver._openai_part_to_gemini_part(
        {
          ["type"] = "text",
          ["bad_text_field"] = "What is in this picture?",
        })

      assert.is_nil(gemini_prompt)
      assert.not_nil(err)

      assert.same("message part type is 'text' but is missing .text block", err)
    end)

    it("transforms an image_url type prompt when data is a URL to gemini GOOD", function()
      local gemini_prompt, err = gemini_driver._openai_part_to_gemini_part(
        {
          ["type"] = "image_url",
          ["image_url"] = {
            ["url"] = "https://example.local/image.jpg",
          },
        })

      assert.not_nil(gemini_prompt)
      assert.is_nil(err)

      assert.same(gemini_prompt,
        {
          ["fileData"] = {
            ["fileUri"] = "https://example.local/image.jpg",
            ["mimeType"] = "image/generic",
          },
        })
    end)

    it("transforms an image_url type prompt when data is a URL to gemini BAD MISSING IMAGE FIELD", function()
      local gemini_prompt, err = gemini_driver._openai_part_to_gemini_part(
        {
          ["type"] = "image_url",
          ["image_url"] = "https://example.local/image.jpg",
        })

      assert.is_nil(gemini_prompt)
      assert.not_nil(err)

      assert.same("message part type is 'image_url' but is missing .image_url.url block", err)
    end)

    it("fails to transform a non-mapped multimodal entity type", function()
      local gemini_prompt, err = gemini_driver._openai_part_to_gemini_part(
        {
          ["type"] = "doesnt_exist",
          ["doesnt_exist"] = "https://example.local/video.mp4",
        })

      assert.is_nil(gemini_prompt)
      assert.not_nil(err)

      assert.same("cannot transform part of type 'doesnt_exist' to Gemini format", err)
    end)

    it("transforms 'describe this image' via URL from openai to gemini", function()
      local gemini_prompt, _, err = gemini_driver._to_gemini_chat_openai(SAMPLE_LLM_V2_CHAT_MULTIMODAL_IMAGE_URL)

      assert.is_nil(err)
      assert.not_nil(gemini_prompt)

      gemini_prompt.generationConfig = nil  -- not needed for comparison

      assert.same({
        ["contents"] = {
          {
            ["role"] = "user",
            ["parts"] = {
              {
                ["text"] = "What is in this picture?",
              },
              {
                ["fileData"] = {
                  ["fileUri"] = "https://example.local/image.jpg",
                  ["mimeType"] = "image/generic",
                },
              }
            },
          },
          {
            ["role"] = "model",
            ["parts"] = {
              {
                ["text"] = "A picture of a cat.",
              },
            },
          },
          {
            ["role"] = "user",
            ["parts"] = {
              {
                ["text"] = "Now draw it wearing a party-hat.",
              },
            },
          },
        }
      }, gemini_prompt)
    end)

    it("transforms 'describe this image' via base64 from openai to gemini", function()
      local gemini_prompt, _, err = gemini_driver._to_gemini_chat_openai(SAMPLE_LLM_V2_CHAT_MULTIMODAL_IMAGE_B64)

      assert.is_nil(err)
      assert.not_nil(gemini_prompt)

      gemini_prompt.generationConfig = nil  -- not needed for comparison

      assert.same({
        ["contents"] = {
          {
            ["role"] = "user",
            ["parts"] = {
              {
                ["text"] = "What is in this picture?",
              },
              {
                ["inlineData"] = {
                  ["data"] = "Y2F0X3BuZ19oZXJlX2xvbAo=",
                  ["mimeType"] = "image/png",
                },
              }
            },
          },
          {
            ["role"] = "model",
            ["parts"] = {
              {
                ["text"] = "A picture of a cat.",
              },
            },
          },
          {
            ["role"] = "user",
            ["parts"] = {
              {
                ["text"] = "Now draw it wearing a party-hat.",
              },
            },
          },
        }
      }, gemini_prompt)
    end)

  end)


  describe("gemini tools", function()
    local gemini_driver

    setup(function()
      _G._TEST = true
      package.loaded["kong.llm.drivers.gemini"] = nil
      gemini_driver = require("kong.llm.drivers.gemini")
    end)

    teardown(function()
      _G._TEST = nil
    end)

    it("transforms openai tools to gemini tools GOOD", function()
      local gemini_tools = gemini_driver._to_tools(SAMPLE_OPENAI_TOOLS_REQUEST.tools)

      assert.not_nil(gemini_tools)
      assert.same(gemini_tools, {
        {
          function_declarations = {
            {
              description = "Check a product is in stock.",
              name = "check_stock",
              parameters = {
                properties = {
                  product_name = {
                    type = "string"
                  }
                },
                required = {
                  "product_name"
                },
                type = "object"
              }
            }
          }
        }
      })
    end)

    it("transforms openai tools to gemini tools NO_TOOLS", function()
      local gemini_tools = gemini_driver._to_tools(SAMPLE_LLM_V1_CHAT)

      assert.is_nil(gemini_tools)
    end)

    it("transforms openai tools to gemini tools NIL", function()
      local gemini_tools = gemini_driver._to_tools(nil)

      assert.is_nil(gemini_tools)
    end)

    it("transforms gemini tools to openai tools GOOD", function()
      local openai_tools = gemini_driver._from_gemini_chat_openai(SAMPLE_GEMINI_TOOLS_RESPONSE, {}, "llm/v1/chat")

      assert.not_nil(openai_tools)

      openai_tools = cjson.decode(openai_tools)
      assert.same(openai_tools.choices[1].message.tool_calls[1]['function'], {
        name = "sql_execute",
        arguments = "{\"product_name\":\"NewPhone\"}"
      })
    end)
  end)

  describe("bedrock tools", function()
    local bedrock_driver

    setup(function()
      _G._TEST = true
      package.loaded["kong.llm.drivers.bedrock"] = nil
      bedrock_driver = require("kong.llm.drivers.bedrock")
    end)

    teardown(function()
      _G._TEST = nil
    end)

    it("transforms openai tools to bedrock tools GOOD", function()
      local bedrock_tools = bedrock_driver._to_tools(SAMPLE_OPENAI_TOOLS_REQUEST.tools)

      assert.not_nil(bedrock_tools)
      assert.same(bedrock_tools, {
        {
          toolSpec = {
            description = "Check a product is in stock.",
            inputSchema = {
              json = {
                properties = {
                  product_name = {
                    type = "string"
                  }
                },
                required = {
                  "product_name"
                },
                type = "object"
              }
            },
            name = "check_stock"
          }
        }
      })
    end)

    it("transforms openai tools to bedrock tools NO_TOOLS", function()
      local bedrock_tools = bedrock_driver._to_tools(SAMPLE_LLM_V1_CHAT)

      assert.is_nil(bedrock_tools)
    end)

    it("transforms openai tools to bedrock tools NIL", function()
      local bedrock_tools = bedrock_driver._to_tools(nil)

      assert.is_nil(bedrock_tools)
    end)

    it("transforms bedrock tools to openai tools GOOD", function()
      local openai_tools = bedrock_driver._from_tool_call_response(SAMPLE_BEDROCK_TOOLS_RESPONSE.output.message.content)

      assert.not_nil(openai_tools)

      assert.same(openai_tools[1]['function'], {
        name = "sumArea",
        arguments = "{\"areas\":[121,212,313]}"
      })
    end)

    it("transforms guardrails into bedrock generation config", function()
      local model_info = {
        route_type = "llm/v1/chat",
        name = "some-model",
        provider = "bedrock",
      }
      local bedrock_guardrails = bedrock_driver._to_bedrock_chat_openai(SAMPLE_LLM_V1_CHAT_WITH_GUARDRAILS, model_info, "llm/v1/chat")

      assert.not_nil(bedrock_guardrails)

      assert.same(bedrock_guardrails.guardrailConfig, {
        ['guardrailIdentifier'] = 'yu5xwvfp4sud',
        ['guardrailVersion'] = '1',
        ['trace'] = 'enabled',
      })
    end)
  end)
end)

describe("json_array_iterator", function()
  local json_array_iterator
  lazy_setup(function()
    _G.TEST = true
    package.loaded["kong.llm.drivers.shared"] = nil
    json_array_iterator = require("kong.llm.drivers.shared")._json_array_iterator
  end)

  -- Helper function to collect all elements from iterator
  local function collect_elements(input)
    local elements = {}
    local iter = json_array_iterator(input)
    local next_element = iter()
    while next_element do
      table.insert(elements, next_element)
      next_element = iter()
    end
    return elements
  end

  it("#qq should handle simple flat arrays", function()
    local input = '[1, 2, 3, 4]'
    local elements = collect_elements(input)
    assert.are.same({"1", "2", "3", "4"}, elements)
  end)

  it("should handle arrays with strings", function()
    local input = '["hello", "world"]'
    local elements = collect_elements(input)
    assert.are.same({'"hello"', '"world"'}, elements)
  end)

  it("should handle nested objects", function()
    local input = '[{"name": "John"}, {"name": "Jane"}]'
    local elements = collect_elements(input)
    assert.are.same(
      '{\"name\": \"John\"}',
      elements[1]
    )
    assert.are.same(
      '{\"name\": \"Jane\"}',
      elements[2]
    )
  end)

  it("should handle nested arrays", function()
    local input = '[[1, 2], [3, 4]]'
    local elements = collect_elements(input)
    assert.are.same({"[1, 2]", "[3, 4]"}, elements)
  end)

  it("should handle whitespace", function()
    local input = '  [  1  ,  2  ]  '
    local elements = collect_elements(input)
    assert.are.same({"1", "2"}, elements)
  end)

  it("should handle empty arrays", function()
    local input = '[]'
    local elements = collect_elements(input)
    assert.are.same({}, elements)
  end)

  it("should handle strings with special characters", function()
    local input = '["{\"special\": \"\\\"quoted\\\"\"}", "[1,2]"]'
    local elements = collect_elements(input)
    assert.are.same(
      {'\"{\"special\": \"\\\"quoted\\\"\"}\"', '"[1,2]"'},
      elements
    )
  end)

  describe("incremental parsing", function()
    it("should handle split within string", function()
      local state = {
        started = false,
        pos = 1,
        input = '',
        eof = false,
      }
      local iter

      -- First chunk (split within string)
      iter = json_array_iterator('["hel', state)
      local element, new_state = iter()
      assert.is_nil(element)  -- Should return nil as string is incomplete
      state = new_state

      -- Second chunk (complete string)
      iter = json_array_iterator('lo"]', state)
      element = iter()
      assert.are.same('"hello"', element)
    end)

    it("should handle split within escaped characters", function()
      local state = {
        started = false,
        pos = 1,
        input = '',
        eof = false,
      }
      local iter

      -- Split during escape sequence
      iter = json_array_iterator('["he\\', state)
      local element, new_state = iter()
      assert.is_nil(element)
      state = new_state

      iter = json_array_iterator('\\nllo"]', state)
      element = iter()
      assert.are.same('"he\\\\nllo"', element)
    end)

    it("should handle split between object braces", function()
      local state = {
        started = false,
        pos = 1,
        input = '',
        eof = false,
      }
      local iter

      -- Split between object definition
      iter = json_array_iterator('[{"name": "Jo', state)
      local element, new_state = iter()
      assert.is_nil(element)
      state = new_state

      iter = json_array_iterator('hn"}, {"age": 30}]', state)
      element = iter()
      assert.are.same('{"name": "John"}', element)

      element = iter()
      assert.are.same('{"age": 30}', element)
    end)

    it("should handle split between array brackets", function()
      local state = {
        started = false,
        pos = 1,
        input = '',
        eof = false,
      }
      local iter

      -- Split between nested array
      iter = json_array_iterator('[[1, 2', state)
      local element, new_state = iter()
      assert.is_nil(element)
      state = new_state

      iter = json_array_iterator('], [3, 4]]', state)
      element = iter()
      assert.are.same('[1, 2]', element)

      element = iter()
      assert.are.same('[3, 4]', element)
    end)

    it("should handle split at comma", function()
      local state = {
        started = false,
        pos = 1,
        input = '',
        eof = false,
      }
      local iter

      -- Split at comma
      iter = json_array_iterator('[1,', state)
      local element, new_state = iter()
      assert.are.same('1', element)
      state = new_state

      iter = json_array_iterator(' 2]', state)
      element = iter()
      assert.are.same('2', element)
    end)

    it("should not split between literal comma", function()
      local state = {
        started = false,
        pos = 1,
        input = '',
        eof = false,
      }
      local iter

      -- Split at comma
      iter = json_array_iterator('[{"message":"hello world"}, {"message":"goodbye,', state)
      local element, _ = iter()
      assert.are.same('{"message":"hello world"}',element)
      local element, _ = iter()
      assert.is_nil(element)

      iter = json_array_iterator(' world"}]', state)
      element = iter()
      assert.are.same('{"message":"goodbye, world"}', element)
    end)

    it("should handle split within complex nested structure", function()
      local state = {
        started = false,
        pos = 1,
        input = '',
        eof = false,
      }
      local iter

      -- Complex nested structure split
      iter = json_array_iterator('[{"users": [{"id": 1', state)
      local element, new_state = iter()
      assert.is_nil(element)
      state = new_state

      iter = json_array_iterator(', "name": "John"}]}, {"status": "', state)
      local element, new_state = iter()
      assert.are.same('{"users": [{"id": 1, "name": "John"}]}', element)
      state = new_state

      iter = json_array_iterator('active"}]', state)
      element = iter()
      assert.are.same('{"status": "active"}', element)
    end)
  end)

  it("should error on invalid start", function()
    assert.has_error(function()
      json_array_iterator('{1, 2, 3}')()
    end, "Invalid start: expected '['")
  end)

  it("should handle complex nested structures", function()
    local input = '[{"users": [{"id": 1, "name": "John"}, {"id": 2, "name": "Jane"}]}, {"status": "active"}]'
    local elements = collect_elements(input)
    assert.are.same(
      '{\"users\": [{\"id\": 1, \"name\": \"John\"}, {\"id\": 2, \"name\": \"Jane\"}]}',
      elements[1]
    )
    assert.are.same(
      '{\"status\": \"active\"}',
      elements[2]
    )
  end)
end)