defmodule Pyex.Stdlib.Hashlib do
  @moduledoc """
  Python `hashlib` module for secure hash and message digest algorithms.

  Provides constructor functions `hashlib.sha256()`, `hashlib.sha1()`,
  `hashlib.md5()`, `hashlib.sha224()`, `hashlib.sha384()`, `hashlib.sha512()`,
  and the generic `hashlib.new(name)` constructor.

  Each returns a hash object with `.update(data)`, `.hexdigest()`, `.digest()`,
  and `.digest_size` / `.block_size` / `.name` attributes.

  Backed by Erlang's `:crypto` module.
  """

  @behaviour Pyex.Stdlib.Module

  @algorithms %{
    "md5" => :md5,
    "sha1" => :sha,
    "sha224" => :sha224,
    "sha256" => :sha256,
    "sha384" => :sha384,
    "sha512" => :sha512
  }

  @digest_sizes %{
    "md5" => 16,
    "sha1" => 20,
    "sha224" => 28,
    "sha256" => 32,
    "sha384" => 48,
    "sha512" => 64
  }

  @block_sizes %{
    "md5" => 64,
    "sha1" => 64,
    "sha224" => 64,
    "sha256" => 64,
    "sha384" => 128,
    "sha512" => 128
  }

  @doc """
  Returns the module value map with hash constructor functions.
  """
  @impl Pyex.Stdlib.Module
  @spec module_value() :: Pyex.Stdlib.Module.module_value()
  def module_value do
    %{
      "md5" => {:builtin, &do_md5/1},
      "sha1" => {:builtin, &do_sha1/1},
      "sha224" => {:builtin, &do_sha224/1},
      "sha256" => {:builtin, &do_sha256/1},
      "sha384" => {:builtin, &do_sha384/1},
      "sha512" => {:builtin, &do_sha512/1},
      "new" => {:builtin, &do_new/1}
    }
  end

  @spec do_md5([Pyex.Interpreter.pyvalue()]) :: Pyex.Interpreter.pyvalue()
  defp do_md5(args), do: hash_with("md5", args)

  @spec do_sha1([Pyex.Interpreter.pyvalue()]) :: Pyex.Interpreter.pyvalue()
  defp do_sha1(args), do: hash_with("sha1", args)

  @spec do_sha224([Pyex.Interpreter.pyvalue()]) :: Pyex.Interpreter.pyvalue()
  defp do_sha224(args), do: hash_with("sha224", args)

  @spec do_sha256([Pyex.Interpreter.pyvalue()]) :: Pyex.Interpreter.pyvalue()
  defp do_sha256(args), do: hash_with("sha256", args)

  @spec do_sha384([Pyex.Interpreter.pyvalue()]) :: Pyex.Interpreter.pyvalue()
  defp do_sha384(args), do: hash_with("sha384", args)

  @spec do_sha512([Pyex.Interpreter.pyvalue()]) :: Pyex.Interpreter.pyvalue()
  defp do_sha512(args), do: hash_with("sha512", args)

  @spec hash_with(String.t(), [Pyex.Interpreter.pyvalue()]) :: Pyex.Interpreter.pyvalue()
  defp hash_with(algo, []), do: make_hash_object(algo, <<>>)
  defp hash_with(algo, [{:bytes, data}]), do: make_hash_object(algo, data)
  defp hash_with(algo, [{:bytearray, data}]), do: make_hash_object(algo, data)

  # CPython raises TypeError for strings; Pyex accepts them for
  # convenience by treating them as UTF-8 bytes.
  defp hash_with(algo, [data]) when is_binary(data),
    do: make_hash_object(algo, data)

  defp hash_with(algo, _),
    do: {:exception, "TypeError: #{algo}() argument must be bytes"}

  @spec do_new([Pyex.Interpreter.pyvalue()]) :: Pyex.Interpreter.pyvalue()
  defp do_new([name]) when is_binary(name) do
    algo = String.downcase(name) |> String.replace("-", "")

    if Map.has_key?(@algorithms, algo) do
      make_hash_object(algo, <<>>)
    else
      {:exception, "ValueError: unsupported hash type #{name}"}
    end
  end

  defp do_new([name, data]) when is_binary(name) and is_binary(data) do
    algo = String.downcase(name) |> String.replace("-", "")

    if Map.has_key?(@algorithms, algo) do
      make_hash_object(algo, data)
    else
      {:exception, "ValueError: unsupported hash type #{name}"}
    end
  end

  defp do_new(_), do: {:exception, "TypeError: new() requires a string algorithm name"}

  @spec make_hash_object(String.t(), binary()) :: Pyex.Interpreter.pyvalue()
  defp make_hash_object(algo_name, data) do
    erlang_algo = Map.get(@algorithms, algo_name)
    digest = :crypto.hash(erlang_algo, data)
    hex = Base.encode16(digest, case: :lower)

    digest_size = Map.get(@digest_sizes, algo_name, 0)
    block_size = Map.get(@block_sizes, algo_name, 0)

    update_fn =
      {:builtin,
       fn
         [_self, {:bytes, new_data}] ->
           {:mutate, make_hash_object(algo_name, data <> new_data), nil}

         [_self, {:bytearray, new_data}] ->
           {:mutate, make_hash_object(algo_name, data <> new_data), nil}

         [_self, new_data] when is_binary(new_data) ->
           {:mutate, make_hash_object(algo_name, data <> new_data), nil}

         [{:bytes, new_data}] ->
           make_hash_object(algo_name, data <> new_data)

         [{:bytearray, new_data}] ->
           make_hash_object(algo_name, data <> new_data)

         [new_data] when is_binary(new_data) ->
           make_hash_object(algo_name, data <> new_data)

         _ ->
           {:exception, "TypeError: update() argument must be bytes"}
       end}

    hexdigest_fn =
      {:builtin,
       fn
         [_self] -> hex
         [] -> hex
       end}

    digest_fn =
      {:builtin,
       fn
         [_self] -> digest
         [] -> digest
       end}

    copy_fn =
      {:builtin,
       fn
         [_self] -> make_hash_object(algo_name, data)
         [] -> make_hash_object(algo_name, data)
       end}

    hash_class =
      {:class, "_hashlib.HASH", [], %{}}

    {:instance, hash_class,
     %{
       "name" => algo_name,
       "digest_size" => digest_size,
       "block_size" => block_size,
       "update" => update_fn,
       "hexdigest" => hexdigest_fn,
       "digest" => digest_fn,
       "copy" => copy_fn
     }}
  end
end
