openai_responses

defmodule OpenAI.Responses.Schema do
  @moduledoc """
  Helper module for defining structured output schemas and function calling tools.

  Converts simple Elixir syntax into JSON Schema format for structured outputs and function parameters.

  ## Array Support (New in 0.6.0)

  Arrays can now be used at the root level of schema definitions. The library automatically
  handles OpenAI's requirement that the root level must be an object by wrapping arrays
  in a temporary object structure and unwrapping them in the response.

  ## Examples

  ### Structured Output Schema with Object

      iex> Responses.Schema.build_output(%{
      ...>   name: {:string, description: "The name of the user"},
      ...>   username: {:string, description: "The username of the user. Must start with @", pattern: "^@[a-zA-Z0-9_]+$"},
      ...>   email: {:string, description: "The email of the user", format: "email"}
      ...> })
      %{
        "name" => "data",
        "type" => "json_schema",
        "strict" => true,
        "schema" => %{
          "type" => "object",
          "properties" => %{
            "name" => %{
              "type" => "string",
              "description" => "The name of the user"
            },
            "username" => %{
              "type" => "string",
              "description" => "The username of the user. Must start with @",
              "pattern" => "^@[a-zA-Z0-9_]+$"
            },
            "email" => %{
              "type" => "string",
              "description" => "The email of the user",
              "format" => "email"
            }
          },
          "additionalProperties" => false,
          "required" => ["name", "username", "email"]
        }
      }

  ### Structured Output Schema with Array at Root

      iex> Responses.Schema.build_output({:array, %{
      ...>   title: :string,
      ...>   completed: :boolean,
      ...>   priority: {:integer, minimum: 1, maximum: 5}
      ...> }})
      %{
        "name" => "data",
        "type" => "json_schema",
        "strict" => true,
        "schema" => %{
          "type" => "object",
          "properties" => %{
            "items" => %{
              "type" => "array",
              "items" => %{
                "type" => "object",
                "properties" => %{
                  "title" => %{"type" => "string"},
                  "completed" => %{"type" => "boolean"},
                  "priority" => %{"type" => "integer", "minimum" => 1, "maximum" => 5}
                },
                "additionalProperties" => false,
                "required" => ["completed", "priority", "title"]
              }
            }
          },
          "additionalProperties" => false,
          "required" => ["items"]
        }
      }

  When using array schemas, the response will be automatically unwrapped so that
  `response.parsed` contains the array directly, not wrapped in an object.

  ### Function Calling Tool

      iex> Responses.Schema.build_function("get_weather", "Get current temperature for a given location.", %{
      ...>   location: {:string, description: "City and country e.g. Bogotá, Colombia"}
      ...> })
      %{
        "name" => "get_weather",
        "type" => "function",
        "strict" => true,
        "description" => "Get current temperature for a given location.",
        "parameters" => %{
          "type" => "object",
          "properties" => %{
            "location" => %{
              "type" => "string",
              "description" => "City and country e.g. Bogotá, Colombia"
            }
          },
          "additionalProperties" => false,
          "required" => ["location"]
        }
      }
  """

  @doc """
  Builds a structured output schema from a simple Elixir map or keyword list format.

  The input should be a map or keyword list where:
  - Keys are field names (atoms)
  - Values are either:
    - A single atom like `:string`, `:number`, `:boolean`, etc.
    - A tuple like `{:string, description: "...", pattern: "..."}`
    - For arrays: `{:array, :string}` or `{:array, %{field: :type}}`

  When using keyword lists, the order of fields is preserved in the required array.
  When using maps, fields are sorted alphabetically in the required array.

  If the root schema is an array, it will be automatically wrapped in an object
  to comply with OpenAI's Structured Outputs requirements.
  """
  def build_output(fields) do
    # Handle array specs at the root level
    schema =
      if array_spec?(fields) do
        build_property(fields)
      else
        build_schema(fields)
      end

    # Check if the root is an array and wrap it if necessary
    schema =
      case schema do
        %{"type" => "array"} = array_schema ->
          %{
            "type" => "object",
            "properties" => %{
              "items" => array_schema
            },
            "additionalProperties" => false,
            "required" => ["items"]
          }

        other ->
          other
      end

    %{
      "name" => "data",
      "type" => "json_schema",
      "strict" => true,
      "schema" => schema
    }
  end

  @doc """
  Builds a function calling tool schema.

  ## Parameters
    - `name` - The function name
    - `description` - A description of what the function does
    - `parameters` - A map or keyword list of parameter definitions (same format as `build_output/1`)

  ## Example

      iex> build_function("get_weather", "Get weather for a location", %{
      ...>   location: {:string, description: "City name"},
      ...>   units: {:string, enum: ["celsius", "fahrenheit"], description: "Temperature units"}
      ...> })
  """
  def build_function(name, description, parameters) do
    %{
      "name" => name,
      "type" => "function",
      "strict" => true,
      "description" => description,
      "parameters" => build_schema(parameters)
    }
  end

  defp build_schema(fields) when is_map(fields) or is_list(fields) do
    build_property(fields)
  end

  defp build_property(spec) do
    spec
    |> normalize_spec()
    |> build_from_normalized()
  end

  # Normalize various input formats to a standard map format
  defp normalize_spec(spec) do
    cond do
      # Simple types
      simple_type?(spec) ->
        normalize_simple_type(spec)

      # Arrays
      array_spec?(spec) ->
        normalize_array_spec(spec)

      # Union types (anyOf)
      anyof_spec?(spec) ->
        normalize_anyof_spec(spec)

      # Type with options
      type_with_options?(spec) ->
        normalize_type_with_options_spec(spec)

      # Lists (keyword lists or empty)
      is_list(spec) ->
        normalize_list_spec(spec)

      # Maps are object specifications
      is_map(spec) ->
        normalize_map_spec(spec)

      # Fallback
      true ->
        raise ArgumentError, "Unsupported schema specification: #{inspect(spec)}"
    end
  end

  # Type checking helpers
  defp simple_type?(type) when is_atom(type) or is_binary(type), do: true
  defp simple_type?(_), do: false

  defp array_spec?({type, _}) when type in [:array, "array"], do: true
  defp array_spec?([type, _]) when type in [:array, "array"], do: true
  defp array_spec?(_), do: false

  defp anyof_spec?({type, specs}) when type in [:anyOf, "anyOf"] and is_list(specs), do: true
  defp anyof_spec?([type, specs]) when type in [:anyOf, "anyOf"] and is_list(specs), do: true
  defp anyof_spec?(_), do: false

  defp type_with_options?([type, opts])
       when (is_atom(type) or is_binary(type)) and (is_list(opts) or is_map(opts)),
       do: true

  defp type_with_options?({type, opts}) when (is_atom(type) or is_binary(type)) and is_list(opts),
    do: true

  defp type_with_options?(_), do: false

  # Normalization helpers
  defp normalize_simple_type(type) do
    %{"type" => to_string(type)}
  end

  defp normalize_array_spec({_array_type, item_spec}) do
    %{"type" => "array", "items" => normalize_spec(item_spec)}
  end

  defp normalize_array_spec([_array_type, item_spec]) do
    %{"type" => "array", "items" => normalize_spec(item_spec)}
  end

  defp normalize_anyof_spec({_anyof_type, specs}) do
    %{"anyOf" => Enum.map(specs, &normalize_spec/1)}
  end

  defp normalize_anyof_spec([_anyof_type, specs]) do
    %{"anyOf" => Enum.map(specs, &normalize_spec/1)}
  end

  defp normalize_type_with_options_spec([type, opts]) do
    opts = normalize_nested_options(opts)
    handle_type_with_options(type, opts, &Map.has_key?/2, &Map.get/2)
  end

  defp normalize_type_with_options_spec({type, opts}) do
    handle_type_with_options(type, opts, &Keyword.has_key?/2, &Keyword.get/2)
  end

  defp handle_type_with_options(type, opts, has_key_fn, get_fn) do
    if type in [:object, "object"] and has_key_fn.(opts, :properties) do
      properties = get_fn.(opts, :properties)
      normalize_object_spec(Enum.to_list(properties))
    else
      normalize_type_with_options(type, opts)
    end
  end

  defp normalize_list_spec([]) do
    normalize_object_spec([])
  end

  defp normalize_list_spec(list) when is_tuple(hd(list)) do
    normalize_object_spec(list)
  end

  defp normalize_list_spec(spec) do
    raise ArgumentError, "Unsupported schema specification: #{inspect(spec)}"
  end

  defp normalize_map_spec(map) do
    map
    |> Enum.sort_by(fn {key, _} -> to_string(key) end)
    |> normalize_object_spec()
  end

  # Handle nested list options like [:max_items, 2] -> [max_items: 2]
  defp normalize_nested_options(opts) when is_list(opts) do
    case opts do
      # Single nested list like [:max_items, 2]
      [[key, value]] when is_atom(key) or is_binary(key) ->
        [{key, value}]

      # Direct list like [:max_items, 2] that's not a keyword list
      [key, value] when (is_atom(key) or is_binary(key)) and not is_tuple(value) ->
        [{key, value}]

      # Already a proper keyword list or other format
      _ ->
        opts
    end
  end

  defp normalize_nested_options(opts), do: opts

  defp normalize_type_with_options(type, opts) do
    base = %{"type" => to_string(type)}

    # Convert options to map if they're a keyword list
    opts_map = if is_list(opts), do: Map.new(opts), else: opts

    # Merge options into base, converting keys to strings
    Enum.reduce(opts_map, base, fn {key, value}, acc ->
      Map.put(acc, to_string(key), value)
    end)
  end

  defp normalize_object_spec(spec) do
    properties =
      spec
      |> Enum.map(fn {name, child_spec} ->
        {to_string(name), normalize_spec(child_spec)}
      end)
      |> Map.new()

    required =
      spec
      |> Enum.map(fn {key, _} -> to_string(key) end)

    %{
      "type" => "object",
      "properties" => properties,
      "required" => required
    }
  end

  # Build the final property from normalized format
  defp build_from_normalized(%{"type" => "array", "items" => items}) do
    %{
      "type" => "array",
      "items" => build_from_normalized(items)
    }
  end

  defp build_from_normalized(%{"anyOf" => specs}) do
    %{
      "anyOf" => Enum.map(specs, &build_from_normalized/1)
    }
  end

  defp build_from_normalized(%{
         "type" => "object",
         "properties" => properties,
         "required" => required
       }) do
    built_properties =
      properties
      |> Enum.map(fn {name, prop} -> {name, build_from_normalized(prop)} end)
      |> Map.new()

    %{
      "type" => "object",
      "properties" => built_properties,
      "additionalProperties" => false,
      "required" => required
    }
  end

  defp build_from_normalized(spec) when is_map(spec) do
    # Already in final format
    spec
  end
end