pyex

defmodule Pyex.Stdlib.YamlParser do
  @moduledoc """
  Hand-rolled block-YAML parser covering the subset used for frontmatter and
  structured config documents generated by LLMs.

  Supports:
  - Block mappings: `key: scalar` and `key:` followed by indented block
  - Block sequences: `- scalar` and `- ` followed by indented block mapping
  - Scalars: strings (plain, single-quoted, double-quoted), integers (decimal,
    octal `0o…`, hex `0x…`), floats (decimal, `1.0e3`, `.inf`, `-.inf`, `.nan`),
    booleans (`true`/`True`/`TRUE`/`false`/`False`/`FALSE`), nulls
    (`null`/`Null`/`NULL`/`~`/empty)
  - `#` comments (full-line and inline, outside quoted strings)
  - Block literal scalars (`|`) and block folded scalars (`>`)

  Does NOT support:
  - Flow style (`{…}`, `[…]`)
  - Anchors, aliases, tags
  - Multi-document streams (`---` / `...` separators)
  - Complex keys

  Returns `{:ok, value}` or `{:error, reason}`.
  """

  @type yaml_value ::
          nil
          | boolean()
          | integer()
          | float()
          | String.t()
          | [yaml_value()]
          | %{String.t() => yaml_value()}

  @max_depth 100

  @doc """
  Parse a YAML string into an Elixir value.

  Keys are always returned as strings.
  """
  @spec parse(String.t()) :: {:ok, yaml_value()} | {:error, String.t()}
  def parse(text) when is_binary(text) do
    case split_lines(text) do
      [] ->
        {:ok, nil}

      [line | _] = lines ->
        col = indent_of(line)

        cond do
          sequence_line?(line, col) -> parse_sequence(lines, col, @max_depth)
          mapping_line?(line, col) -> parse_mapping(lines, col, @max_depth)
          true -> {:error, "unexpected line: #{inspect(line)}"}
        end
    end
  end

  @spec split_lines(String.t()) :: [String.t()]
  defp split_lines(text) do
    text
    |> String.split("\n")
    |> Enum.map(&strip_inline_comment/1)
    |> Enum.reject(&blank?/1)
  end

  # ---------------------------------------------------------------------------
  # Mapping
  # ---------------------------------------------------------------------------

  @spec parse_mapping([String.t()], non_neg_integer(), non_neg_integer()) ::
          {:ok, %{String.t() => yaml_value()}} | {:error, String.t()}
  defp parse_mapping(lines, indent, depth) do
    do_mapping(lines, indent, depth, %{})
  end

  @spec do_mapping([String.t()], non_neg_integer(), non_neg_integer(), map()) ::
          {:ok, map()} | {:error, String.t()}
  defp do_mapping([], _indent, _depth, acc), do: {:ok, acc}

  defp do_mapping([line | _rest] = lines, indent, depth, acc) do
    col = indent_of(line)

    cond do
      col < indent ->
        {:ok, acc}

      col > indent ->
        {:error, "unexpected indent on line: #{inspect(line)}"}

      mapping_line?(line, col) ->
        parse_mapping_entry(lines, indent, depth, acc)

      true ->
        {:ok, acc}
    end
  end

  @spec parse_mapping_entry([String.t()], non_neg_integer(), non_neg_integer(), map()) ::
          {:ok, map()} | {:error, String.t()}
  defp parse_mapping_entry([line | rest], indent, depth, acc) do
    trimmed = String.slice(line, indent, byte_size(line))

    case split_mapping_key(trimmed) do
      {:ok, key, value_part} ->
        if depth == 0 do
          {:error, "maximum nesting depth exceeded"}
        else
          case parse_mapping_value(value_part, rest, indent, depth - 1) do
            {:ok, value, remaining} ->
              do_mapping(remaining, indent, depth, Map.put(acc, key, value))

            {:error, _} = err ->
              err
          end
        end

      :error ->
        {:error, "invalid mapping key on line: #{inspect(line)}"}
    end
  end

  @spec split_mapping_key(String.t()) :: {:ok, String.t(), String.t()} | :error
  defp split_mapping_key(line) do
    case parse_key(line) do
      {:ok, key, rest} ->
        case rest do
          ":" <> after_colon -> {:ok, key, String.trim_leading(after_colon, " ")}
          _ -> :error
        end

      :error ->
        :error
    end
  end

  @spec parse_key(String.t()) :: {:ok, String.t(), String.t()} | :error
  defp parse_key(<<"\"", rest::binary>>) do
    case parse_double_quoted(rest) do
      {:ok, key, after_quote} -> {:ok, key, after_quote}
      :error -> :error
    end
  end

  defp parse_key(<<"'", rest::binary>>) do
    case parse_single_quoted(rest) do
      {:ok, key, after_quote} -> {:ok, key, after_quote}
      :error -> :error
    end
  end

  defp parse_key(line) do
    case :binary.match(line, ":") do
      {idx, _} ->
        key = binary_part(line, 0, idx)
        rest = binary_part(line, idx, byte_size(line) - idx)

        if plain_key?(key) do
          {:ok, String.trim(key), rest}
        else
          :error
        end

      :nomatch ->
        :error
    end
  end

  @spec plain_key?(String.t()) :: boolean()
  defp plain_key?(s), do: s =~ ~r/\A[^\[\]{},#|>&*!'"@`%][^\[\]{},#]*\z/

  @spec parse_mapping_value(String.t(), [String.t()], non_neg_integer(), non_neg_integer()) ::
          {:ok, yaml_value(), [String.t()]} | {:error, String.t()}
  defp parse_mapping_value("", rest, indent, depth) do
    parse_block_value(rest, indent, depth)
  end

  defp parse_mapping_value("|" <> _opts, rest, indent, _depth) do
    parse_block_scalar(rest, indent, :literal)
  end

  defp parse_mapping_value(">" <> _opts, rest, indent, _depth) do
    parse_block_scalar(rest, indent, :folded)
  end

  defp parse_mapping_value("[" <> _, _rest, _indent, _depth) do
    {:error, "flow sequences ([...]) are not supported"}
  end

  defp parse_mapping_value("{" <> _, _rest, _indent, _depth) do
    {:error, "flow mappings ({...}) are not supported"}
  end

  defp parse_mapping_value(value_str, rest, _indent, _depth) do
    {:ok, coerce_scalar(value_str), rest}
  end

  @spec parse_block_value([String.t()], non_neg_integer(), non_neg_integer()) ::
          {:ok, yaml_value(), [String.t()]} | {:error, String.t()}
  defp parse_block_value([], _indent, _depth), do: {:ok, nil, []}

  defp parse_block_value([next | _] = rest, indent, depth) do
    child_col = indent_of(next)

    cond do
      child_col <= indent ->
        {:ok, nil, rest}

      sequence_line?(next, child_col) ->
        case parse_sequence(rest, child_col, depth) do
          {:ok, seq} -> {:ok, seq, drop_while_indented(rest, child_col)}
          {:error, _} = err -> err
        end

      mapping_line?(next, child_col) ->
        case parse_mapping(rest, child_col, depth) do
          {:ok, map} -> {:ok, map, drop_while_indented(rest, child_col)}
          {:error, _} = err -> err
        end

      true ->
        {:ok, nil, rest}
    end
  end

  # ---------------------------------------------------------------------------
  # Sequence
  # ---------------------------------------------------------------------------

  @spec parse_sequence([String.t()], non_neg_integer(), non_neg_integer()) ::
          {:ok, [yaml_value()]} | {:error, String.t()}
  defp parse_sequence(lines, indent, depth) do
    do_sequence(lines, indent, depth, [])
  end

  @spec do_sequence([String.t()], non_neg_integer(), non_neg_integer(), [yaml_value()]) ::
          {:ok, [yaml_value()]} | {:error, String.t()}
  defp do_sequence([], _indent, _depth, acc), do: {:ok, Enum.reverse(acc)}

  defp do_sequence([line | _rest] = lines, indent, depth, acc) do
    col = indent_of(line)

    cond do
      col < indent ->
        {:ok, Enum.reverse(acc)}

      col > indent ->
        {:error, "unexpected indent in sequence: #{inspect(line)}"}

      sequence_line?(line, col) ->
        if depth == 0 do
          {:error, "maximum nesting depth exceeded"}
        else
          parse_sequence_item(lines, indent, depth, acc)
        end

      true ->
        {:ok, Enum.reverse(acc)}
    end
  end

  @spec parse_sequence_item([String.t()], non_neg_integer(), non_neg_integer(), [yaml_value()]) ::
          {:ok, [yaml_value()]} | {:error, String.t()}
  defp parse_sequence_item([line | rest], indent, depth, acc) do
    prefix_len = indent + 2
    value_part = String.slice(line, prefix_len, byte_size(line))

    case parse_sequence_value(value_part, rest, indent, depth - 1) do
      {:ok, value, remaining} ->
        do_sequence(remaining, indent, depth, [value | acc])

      {:error, _} = err ->
        err
    end
  end

  @spec parse_sequence_value(String.t(), [String.t()], non_neg_integer(), non_neg_integer()) ::
          {:ok, yaml_value(), [String.t()]} | {:error, String.t()}
  defp parse_sequence_value("", rest, indent, depth) do
    parse_block_value(rest, indent, depth)
  end

  defp parse_sequence_value(value_str, rest, indent, depth) do
    trimmed = String.trim_leading(value_str)

    cond do
      mapping_line?(trimmed, 0) ->
        inline_line = String.duplicate(" ", indent + 2) <> trimmed

        case parse_mapping([inline_line | rest], indent + 2, depth) do
          {:ok, map} ->
            remaining = drop_while_indented([inline_line | rest], indent + 2)

            {:ok, map, remaining}

          {:error, _} = err ->
            err
        end

      true ->
        {:ok, coerce_scalar(trimmed), rest}
    end
  end

  # ---------------------------------------------------------------------------
  # Block scalars
  # ---------------------------------------------------------------------------

  @spec parse_block_scalar([String.t()], non_neg_integer(), :literal | :folded) ::
          {:ok, String.t(), [String.t()]} | {:error, String.t()}
  defp parse_block_scalar([], _indent, _style), do: {:ok, "", []}

  defp parse_block_scalar([first | _] = lines, indent, style) do
    child_col = indent_of(first)

    if child_col <= indent do
      {:ok, "", lines}
    else
      {block_lines, rest} = Enum.split_while(lines, &(indent_of(&1) >= child_col))

      content =
        block_lines
        |> Enum.map(&String.slice(&1, child_col, byte_size(&1)))

      text =
        case style do
          :literal -> Enum.join(content, "\n") <> "\n"
          :folded -> Enum.join(content, " ") <> "\n"
        end

      {:ok, text, rest}
    end
  end

  # ---------------------------------------------------------------------------
  # Helpers
  # ---------------------------------------------------------------------------

  @spec drop_while_indented([String.t()], non_neg_integer()) :: [String.t()]
  defp drop_while_indented(lines, indent) do
    Enum.drop_while(lines, &(indent_of(&1) >= indent))
  end

  @spec indent_of(String.t()) :: non_neg_integer()
  defp indent_of(line) do
    byte_size(line) - byte_size(String.trim_leading(line, " "))
  end

  @spec blank?(String.t()) :: boolean()
  defp blank?(line), do: String.trim(line) == ""

  @spec sequence_line?(String.t(), non_neg_integer()) :: boolean()
  defp sequence_line?(line, col) do
    String.starts_with?(String.slice(line, col, byte_size(line)), "- ")
  end

  @spec mapping_line?(String.t(), non_neg_integer()) :: boolean()
  defp mapping_line?(line, col) do
    trimmed = String.slice(line, col, byte_size(line))
    not sequence_line?(line, col) and String.contains?(trimmed, ":")
  end

  # ---------------------------------------------------------------------------
  # Scalar coercion (YAML 1.2 core schema)
  # ---------------------------------------------------------------------------

  @spec coerce_scalar(String.t()) :: yaml_value()
  defp coerce_scalar(""), do: nil
  defp coerce_scalar("~"), do: nil
  defp coerce_scalar("null"), do: nil
  defp coerce_scalar("Null"), do: nil
  defp coerce_scalar("NULL"), do: nil
  defp coerce_scalar("true"), do: true
  defp coerce_scalar("True"), do: true
  defp coerce_scalar("TRUE"), do: true
  defp coerce_scalar("false"), do: false
  defp coerce_scalar("False"), do: false
  defp coerce_scalar("FALSE"), do: false
  defp coerce_scalar(".inf"), do: :infinity
  defp coerce_scalar(".Inf"), do: :infinity
  defp coerce_scalar(".INF"), do: :infinity
  defp coerce_scalar("-.inf"), do: :neg_infinity
  defp coerce_scalar("-.Inf"), do: :neg_infinity
  defp coerce_scalar("-.INF"), do: :neg_infinity
  defp coerce_scalar(".nan"), do: :nan
  defp coerce_scalar(".NaN"), do: :nan
  defp coerce_scalar(".NAN"), do: :nan

  defp coerce_scalar(<<"\"", rest::binary>>) do
    case parse_double_quoted(rest) do
      {:ok, s, _} -> s
      :error -> rest
    end
  end

  defp coerce_scalar(<<"'", rest::binary>>) do
    case parse_single_quoted(rest) do
      {:ok, s, _} -> s
      :error -> rest
    end
  end

  defp coerce_scalar(<<"0x", hex::binary>>), do: parse_hex(hex)
  defp coerce_scalar(<<"0X", hex::binary>>), do: parse_hex(hex)
  defp coerce_scalar(<<"0o", oct::binary>>), do: parse_oct(oct)
  defp coerce_scalar(<<"0O", oct::binary>>), do: parse_oct(oct)

  defp coerce_scalar(s) do
    case Integer.parse(s) do
      {n, ""} ->
        n

      _ ->
        case Float.parse(s) do
          {f, ""} -> f
          _ -> s
        end
    end
  end

  @spec parse_hex(String.t()) :: integer() | String.t()
  defp parse_hex(s) do
    case Integer.parse(s, 16) do
      {n, ""} -> n
      _ -> "0x" <> s
    end
  end

  @spec parse_oct(String.t()) :: integer() | String.t()
  defp parse_oct(s) do
    case Integer.parse(s, 8) do
      {n, ""} -> n
      _ -> "0o" <> s
    end
  end

  @spec parse_double_quoted(String.t()) :: {:ok, String.t(), String.t()} | :error
  defp parse_double_quoted(s), do: do_parse_double(s, [])

  @spec do_parse_double(String.t(), iodata()) :: {:ok, String.t(), String.t()} | :error
  defp do_parse_double("", _acc), do: :error
  defp do_parse_double("\"" <> rest, acc), do: {:ok, IO.iodata_to_binary(acc), rest}
  defp do_parse_double("\\n" <> rest, acc), do: do_parse_double(rest, [acc, "\n"])
  defp do_parse_double("\\t" <> rest, acc), do: do_parse_double(rest, [acc, "\t"])
  defp do_parse_double("\\r" <> rest, acc), do: do_parse_double(rest, [acc, "\r"])
  defp do_parse_double("\\\\" <> rest, acc), do: do_parse_double(rest, [acc, "\\"])
  defp do_parse_double("\\\"" <> rest, acc), do: do_parse_double(rest, [acc, "\""])

  defp do_parse_double(<<c::utf8, rest::binary>>, acc) do
    do_parse_double(rest, [acc, <<c::utf8>>])
  end

  @spec parse_single_quoted(String.t()) :: {:ok, String.t(), String.t()} | :error
  defp parse_single_quoted(s), do: do_parse_single(s, [])

  @spec do_parse_single(String.t(), iodata()) :: {:ok, String.t(), String.t()} | :error
  defp do_parse_single("", _acc), do: :error
  defp do_parse_single("''" <> rest, acc), do: do_parse_single(rest, [acc, "'"])
  defp do_parse_single("'" <> rest, acc), do: {:ok, IO.iodata_to_binary(acc), rest}

  defp do_parse_single(<<c::utf8, rest::binary>>, acc) do
    do_parse_single(rest, [acc, <<c::utf8>>])
  end

  @spec strip_inline_comment(String.t()) :: String.t()
  defp strip_inline_comment(line) do
    case find_comment_start(line, false, false) do
      nil -> line
      idx -> String.slice(line, 0, idx) |> String.trim_trailing()
    end
  end

  @spec find_comment_start(String.t(), boolean(), boolean()) :: non_neg_integer() | nil
  defp find_comment_start(<<>>, _in_dq, _in_sq), do: nil

  defp find_comment_start(<<"\\\"", rest::binary>>, true, false) do
    case find_comment_start(rest, true, false) do
      nil -> nil
      n -> n + 2
    end
  end

  defp find_comment_start(<<"\\", _::utf8, rest::binary>>, in_dq, in_sq) when in_dq or in_sq do
    case find_comment_start(rest, in_dq, in_sq) do
      nil -> nil
      n -> n + 2
    end
  end

  defp find_comment_start(<<"\"", rest::binary>>, false, false) do
    case find_comment_start(rest, true, false) do
      nil -> nil
      n -> n + 1
    end
  end

  defp find_comment_start(<<"\"", rest::binary>>, true, false) do
    case find_comment_start(rest, false, false) do
      nil -> nil
      n -> n + 1
    end
  end

  defp find_comment_start(<<"'", rest::binary>>, false, false) do
    case find_comment_start(rest, false, true) do
      nil -> nil
      n -> n + 1
    end
  end

  defp find_comment_start(<<"'", rest::binary>>, false, true) do
    case find_comment_start(rest, false, false) do
      nil -> nil
      n -> n + 1
    end
  end

  defp find_comment_start(<<"#", _::binary>>, false, false) do
    0
  end

  defp find_comment_start(<<_::utf8, rest::binary>>, in_dq, in_sq) do
    case find_comment_start(rest, in_dq, in_sq) do
      nil -> nil
      n -> n + 1
    end
  end
end