defmodule Bumblebee.HuggingFace.HubTest do
  use ExUnit.Case, async: true

  alias Bumblebee.HuggingFace.Hub

  setup do
    bypass = Bypass.open()
    {:ok, bypass: bypass}
  end

  describe "cached_download/2" do
    @tag :tmp_dir
    test "checks etag and downloads the file", %{bypass: bypass, tmp_dir: tmp_dir} do
      Bypass.expect_once(bypass, "HEAD", "/file.json", fn conn ->
        serve_with_etag(conn, ~s/"hash"/, "")
      end)

      Bypass.expect_once(bypass, "GET", "/file.json", fn conn ->
        serve_with_etag(conn, ~s/"hash"/, "{}")
      end)

      url = url(bypass.port) <> "/file.json"

      assert {:ok, path} = Hub.cached_download(url, cache_dir: tmp_dir, offline: false)
      assert File.read!(path) == "{}"
    end

    @tag :tmp_dir
    test "returns a cached file when etag matches", %{bypass: bypass, tmp_dir: tmp_dir} do
      Bypass.expect(bypass, "HEAD", "/file.json", fn conn ->
        serve_with_etag(conn, ~s/"hash"/, "")
      end)

      Bypass.expect_once(bypass, "GET", "/file.json", fn conn ->
        serve_with_etag(conn, ~s/"hash"/, "{}")
      end)

      url = url(bypass.port) <> "/file.json"

      assert {:ok, path} = Hub.cached_download(url, cache_dir: tmp_dir, offline: false)
      assert {:ok, ^path} = Hub.cached_download(url, cache_dir: tmp_dir, offline: false)
      assert File.read!(path) == "{}"
    end

    @tag :tmp_dir
    test "redownloads the file when etag changes", %{bypass: bypass, tmp_dir: tmp_dir} do
      counter = start_supervised!({Agent, fn -> 0 end})

      Bypass.expect(bypass, "HEAD", "/file.json", fn conn ->
        case Agent.get(counter, & &1) do
          0 -> serve_with_etag(conn, ~s/"hash"/, "")
          1 -> serve_with_etag(conn, ~s/"hash2"/, "[]")
        end
      end)

      Bypass.expect(bypass, "GET", "/file.json", fn conn ->
        counter
        |> Agent.get_and_update(fn counter -> {counter, counter + 1} end)
        |> case do
          0 -> serve_with_etag(conn, ~s/"hash"/, "{}")
          1 -> serve_with_etag(conn, ~s/"hash2"/, "[]")
        end
      end)

      url = url(bypass.port) <> "/file.json"

      assert {:ok, path1} = Hub.cached_download(url, cache_dir: tmp_dir, offline: false)
      assert {:ok, path2} = Hub.cached_download(url, cache_dir: tmp_dir, offline: false)
      assert path1 != path2
      assert File.read!(path2) == "[]"
    end

    @tag :tmp_dir
    test "handles redirected downloads", %{bypass: bypass, tmp_dir: tmp_dir} do
      Bypass.expect_once(bypass, "HEAD", "/file.bin", fn conn ->
        url = Plug.Conn.request_url(conn)

        conn
        |> Plug.Conn.put_resp_header("x-linked-etag", ~s/"hash"/)
        |> Plug.Conn.put_resp_header("location", url <> "/storage")
        |> Plug.Conn.resp(302, "")
      end)

      Bypass.expect_once(bypass, "GET", "/file.bin/storage", fn conn ->
        conn
        |> Plug.Conn.put_resp_header("etag", ~s/"hash"/)
        |> Plug.Conn.resp(200, <<0, 1>>)
      end)

      url = url(bypass.port) <> "/file.bin"

      assert {:ok, path} = Hub.cached_download(url, cache_dir: tmp_dir, offline: false)
      assert File.read!(path) == <<0, 1>>
    end

    @tag :tmp_dir
    test "caches redirected downloads based on x-linked-etag", %{bypass: bypass, tmp_dir: tmp_dir} do
      Bypass.expect(bypass, "HEAD", "/file.bin", fn conn ->
        url = Plug.Conn.request_url(conn)

        conn
        |> Plug.Conn.put_resp_header("x-linked-etag", ~s/"hash"/)
        |> Plug.Conn.put_resp_header("location", url <> "/storage")
        |> Plug.Conn.resp(302, "")
      end)

      Bypass.expect_once(bypass, "GET", "/file.bin/storage", fn conn ->
        serve_with_etag(conn, ~s/"hash"/, <<0, 1>>)
      end)

      url = url(bypass.port) <> "/file.bin"

      assert {:ok, path} = Hub.cached_download(url, cache_dir: tmp_dir, offline: false)
      assert {:ok, ^path} = Hub.cached_download(url, cache_dir: tmp_dir, offline: false)
      assert File.read!(path) == <<0, 1>>
    end

    @tag :tmp_dir
    test "follows relative redirect before checking etag", %{bypass: bypass, tmp_dir: tmp_dir} do
      Bypass.expect_once(bypass, "HEAD", "/repo/file.json", fn conn ->
        conn
        |> Plug.Conn.put_resp_header("location", "/repo-renamed/file.json")
        |> Plug.Conn.resp(307, "")
      end)

      Bypass.expect_once(bypass, "HEAD", "/repo-renamed/file.json", fn conn ->
        serve_with_etag(conn, ~s/"hash"/, "{}")
      end)

      Bypass.expect_once(bypass, "GET", "/repo-renamed/file.json", fn conn ->
        serve_with_etag(conn, ~s/"hash"/, "{}")
      end)

      url = url(bypass.port) <> "/repo/file.json"

      assert {:ok, path} = Hub.cached_download(url, cache_dir: tmp_dir, offline: false)
      assert File.read!(path) == "{}"
    end

    @tag :tmp_dir
    test "returns an error on missing etag header", %{bypass: bypass, tmp_dir: tmp_dir} do
      Bypass.expect_once(bypass, "HEAD", "/file.json", fn conn ->
        Plug.Conn.resp(conn, 200, "")
      end)

      url = url(bypass.port) <> "/file.json"

      assert {:error, "no ETag found on the resource"} =
               Hub.cached_download(url, cache_dir: tmp_dir, offline: false)
    end

    @tag :tmp_dir
    test "returns an error on http failure", %{bypass: bypass, tmp_dir: tmp_dir} do
      Bypass.expect_once(bypass, "HEAD", "/file.json", fn conn ->
        Plug.Conn.resp(conn, 500, "")
      end)

      url = url(bypass.port) <> "/file.json"

      assert {:error, "HTTP request failed with status 500, url: " <> _} =
               Hub.cached_download(url, cache_dir: tmp_dir, offline: false)
    end

    @tag :tmp_dir
    test "returns more specific error if x-error-code is present",
         %{bypass: bypass, tmp_dir: tmp_dir} do
      Bypass.expect_once(bypass, "HEAD", "/file.json", fn conn ->
        conn
        |> Plug.Conn.put_resp_header("x-error-code", "RepoNotFound")
        |> Plug.Conn.resp(404, "")
      end)

      url = url(bypass.port) <> "/file.json"

      assert {:error, "repository not found, url: " <> _} =
               Hub.cached_download(url, cache_dir: tmp_dir, offline: false)
    end

    @tag :tmp_dir
    test "returns cached without checking etag when :offline is enabled",
         %{bypass: bypass, tmp_dir: tmp_dir} do
      Bypass.expect_once(bypass, "HEAD", "/file.json", fn conn ->
        serve_with_etag(conn, ~s/"hash"/, "")
      end)

      Bypass.expect_once(bypass, "GET", "/file.json", fn conn ->
        serve_with_etag(conn, ~s/"hash"/, "{}")
      end)

      url = url(bypass.port) <> "/file.json"

      assert {:ok, path} = Hub.cached_download(url, cache_dir: tmp_dir, offline: false)
      assert File.read!(path) == "{}"

      assert {:ok, path} = Hub.cached_download(url, cache_dir: tmp_dir, offline: true)
      assert File.read!(path) == "{}"
    end

    @tag :tmp_dir
    test "returns an error when :offline is enabled and file not in cache",
         %{bypass: bypass, tmp_dir: tmp_dir} do
      url = url(bypass.port) <> "/file.json"

      assert {:error,
              "could not find file in local cache and outgoing traffic is disabled, url: " <> _} =
               Hub.cached_download(url, cache_dir: tmp_dir, offline: true)
    end

    @tag :tmp_dir
    test "includes authorization header when :auth_token is given",
         %{bypass: bypass, tmp_dir: tmp_dir} do
      Bypass.expect_once(bypass, "HEAD", "/file.json", fn conn ->
        assert {"authorization", "Bearer token"} in conn.req_headers

        serve_with_etag(conn, ~s/"hash"/, "")
      end)

      Bypass.expect_once(bypass, "GET", "/file.json", fn conn ->
        assert {"authorization", "Bearer token"} in conn.req_headers

        serve_with_etag(conn, ~s/"hash"/, "{}")
      end)

      url = url(bypass.port) <> "/file.json"

      assert {:ok, _path} =
               Hub.cached_download(url, auth_token: "token", cache_dir: tmp_dir, offline: false)
    end

    @tag :tmp_dir
    test "skips authentication header for redirected requests",
         %{bypass: bypass, tmp_dir: tmp_dir} do
      # Context: HuggingFace Hub returns redirects for files stored
      # in LFS and the redirect location already has S3 signature in
      # URL params. If the location points to S3 directly, which is
      # the case within HF spaces, passing the Authorization header
      # leads to failure, presumably because it takes precedence over
      # the params signature

      Bypass.expect_once(bypass, "HEAD", "/file.bin", fn conn ->
        assert {"authorization", "Bearer token"} in conn.req_headers

        url = Plug.Conn.request_url(conn)

        conn
        |> Plug.Conn.put_resp_header("x-linked-etag", ~s/"hash"/)
        |> Plug.Conn.put_resp_header("location", url <> "/storage")
        |> Plug.Conn.resp(302, "")
      end)

      Bypass.expect_once(bypass, "GET", "/file.bin/storage", fn conn ->
        assert {"authorization", "Bearer token"} not in conn.req_headers

        conn
        |> Plug.Conn.put_resp_header("etag", ~s/"hash"/)
        |> Plug.Conn.resp(200, <<0, 1>>)
      end)

      url = url(bypass.port) <> "/file.bin"

      assert {:ok, _path} =
               Hub.cached_download(url, auth_token: "token", cache_dir: tmp_dir, offline: false)
    end
  end

  defp url(port), do: "http://localhost:#{port}"

  defp serve_with_etag(conn, etag, body) do
    conn
    |> Plug.Conn.put_resp_header("etag", etag)
    |> Plug.Conn.resp(200, body)
  end
end
