Code.require_file("tpl.exs")

defmodule TplProlog do
  # Prolog comparison/arithmetic operators not in the TPL precedence table.
  # These remain in natural order [lhs, op, rhs] in the AST.
  @infix_ops ["\\=", ">=", "=<", ">", "<", "is", "==", "\\==", "=:=", "=\\="]

  # --- Public API ---

  def to_prolog(ast_nodes) do
    ast_nodes
    |> Enum.reject(&skip?/1)
    |> Enum.map(&emit(&1, :top))
    |> Enum.reject(&is_nil/1)
    |> Enum.join("\n")
  end

  # --- Node emitters ---

  # emit/2: convert one AST node to a Prolog string in the given context.
  # context is :top (clause level), :body (goal), or :arg (term inside a call).

  defp emit(%{t: :sexp, list: list}, ctx), do: sexp(clean(list), ctx)
  defp emit(%{t: :brace, list: list}, _ctx), do: body(list)
  defp emit(%{t: :token, str: s}, _ctx), do: s
  defp emit(%{t: :string, chunks: chunks, values: []}, _ctx), do: ~s("#{Enum.join(chunks)}")
  defp emit(%{t: :string, chunks: chunks, values: values}, _ctx) do
    # Best-effort: interleave chunks and values
    max = max(length(chunks), length(values))

    Enum.flat_map(0..(max - 1), fn i ->
      chunk = if i < length(chunks), do: [Enum.at(chunks, i)], else: []
      value = if i < length(values), do: ["#{emit(Enum.at(values, i), :arg)}"], else: []
      chunk ++ value
    end)
    |> Enum.join()
    |> then(&~s("#{&1}"))
  end
  defp emit(%{t: t}, _ctx) when t in [:newline, :comment, :dot, :comma], do: nil

  # --- Helpers ---

  defp skip?(%{t: t}) when t in [:newline, :comment], do: true
  defp skip?(_), do: false

  defp clean(list), do: Enum.reject(list, &skip?/1)

  defp dot(s), do: s <> "."
  defp maybe_dot(s, :top), do: dot(s)
  defp maybe_dot(s, _), do: s

  # term/1: list of head nodes -> Prolog term (functor with args, or atom)
  defp term([%{t: :token, str: f}]), do: f
  defp term([%{t: :token, str: f} | args]) do
    args_s = args |> Enum.map(&emit(&1, :arg)) |> Enum.reject(&is_nil/1) |> Enum.join(", ")
    if args_s == "", do: f, else: "#{f}(#{args_s})"
  end
  defp term([single]), do: emit(single, :arg)
  defp term(nodes), do: nodes |> Enum.map(&emit(&1, :arg)) |> Enum.reject(&is_nil/1) |> Enum.join(", ")

  # body/1: goal list from a brace block -> comma-joined conjunction string
  defp body(list) do
    goals =
      list
      |> Enum.reject(&skip?/1)
      |> Enum.map(&emit(&1, :body))
      |> Enum.reject(&is_nil/1)

    case goals do
      [] -> "true"
      _ -> Enum.join(goals, ",\n    ")
    end
  end

  # findall_body: inline conjunction with parens for multi-goal bodies.
  # Avoids findall(T, goal1,\n    goal2, R) misparse.
  defp findall_body(list) do
    goals = list |> Enum.reject(&skip?/1) |> Enum.map(&emit(&1, :body)) |> Enum.reject(&is_nil/1)
    case goals do
      [] -> "true"
      [single] -> single
      multiple -> "(#{Enum.join(multiple, ", ")})"
    end
  end

  defp brace?(%{t: :brace}), do: true
  defp brace?(_), do: false

  # --- S-expression dispatch ---

  # Any sexp whose last element is a brace is a rule (or DCG rule).
  # e.g.  ancestor X Y { parent X Z, ancestor Z Y }
  #       dcg command (cmd A T) { verb A, noun T }
  defp sexp(list, ctx) when length(list) >= 2 do
    head = Enum.drop(list, -1)
    last = List.last(list)

    # A rule has exactly one trailing brace (the body) and no braces in the head.
    # This excludes `or { A } { B }` where the head itself contains braces.
    if brace?(last) and not Enum.any?(head, &brace?/1) do
      body_s = body(last.list)

      case head do
        [%{t: :token, str: "dcg"} | rest] ->
          "#{term(rest)} --> #{body_s}" |> maybe_dot(ctx)

        _ ->
          "#{term(head)} :- #{body_s}" |> maybe_dot(ctx)
      end
    else
      sexp_no_rule(list, ctx)
    end
  end

  defp sexp(list, ctx), do: sexp_no_rule(list, ctx)

  defp sexp_no_rule(list, ctx) do
    case list do
      # ── Findall ──────────────────────────────────────────────────────────────
      # Var = find Template { Goals }
      # AST (after precedence pass): [token:=, token:Var, sexp:[find, Template, brace:Goals]]
      [%{t: :token, str: "="}, %{t: :token, str: var}, %{t: :sexp, list: rhs}] ->
        rhs_clean = clean(rhs)

        if match?([%{t: :token, str: "find"} | _], rhs_clean) and
             length(rhs_clean) >= 3 and brace?(List.last(rhs_clean)) do
          [_ | rest] = rhs_clean
          {tmpl_nodes, [brace]} = Enum.split(rest, length(rest) - 1)
          tmpl_s = term(tmpl_nodes)
          goal_s = findall_body(brace.list)
          "findall(#{tmpl_s}, #{goal_s}, #{var})" |> maybe_dot(ctx)
        else
          rhs_s = sexp(rhs_clean, :arg)
          "#{var} = #{rhs_s}" |> maybe_dot(ctx)
        end

      # ── Assignment / unification ─────────────────────────────────────────────
      # = lhs rhs  (operator-first because = is in the TPL precedence table)
      [%{t: :token, str: "="}, lhs, rhs] ->
        "#{emit(lhs, :body)} = #{emit(rhs, :body)}" |> maybe_dot(ctx)

      # ── Infix comparison operators (not in TPL table, stay in natural order) ──
      # e.g.  A >= 18  ,  X \= Y
      [lhs, %{t: :token, str: op}, rhs] when op in @infix_ops ->
        "#{emit(lhs, :body)} #{op} #{emit(rhs, :body)}"

      # ── Disjunction ──────────────────────────────────────────────────────────
      # or { goalA } { goalB } ...
      [%{t: :token, str: "or"} | branches] ->
        branches
        |> Enum.filter(&match?(%{t: :brace}, &1))
        |> Enum.map(&emit(&1, :body))
        |> Enum.join(" ; ")
        |> then(&"(#{&1})")

      # ── Negation ─────────────────────────────────────────────────────────────
      # not goal...
      [%{t: :token, str: "not"} | goals] ->
        "\\+ #{term(goals)}"

      # ── General functor call / fact ──────────────────────────────────────────
      [%{t: :token, str: f} | args] ->
        args_s = args |> Enum.map(&emit(&1, :arg)) |> Enum.reject(&is_nil/1) |> Enum.join(", ")
        (if args_s == "", do: f, else: "#{f}(#{args_s})") |> maybe_dot(ctx)

      # ── Fallback ─────────────────────────────────────────────────────────────
      _ ->
        list
        |> Enum.map(&emit(&1, :body))
        |> Enum.reject(&is_nil/1)
        |> Enum.join(" ")
        |> maybe_dot(ctx)
    end
  end
end


# ── Demo ──────────────────────────────────────────────────────────────────────

prolog_code =
  """
  parent elizabeth charles
  parent charles william
  parent charles harry

  male charles
  male william
  male harry
  female elizabeth

  age charles 75
  age william 42
  age harry 40

  person X {
    male X
  }

  grandparent G C {
    parent G P
    parent P C
    (G \= C)
  }

  is_adult Person {
    age Person A
    A >= 18
  }

  ancestor X Y {
    parent X Z
    ancestor Z Y
  }

  child_or_pet_owner X {
    or {
      parent X _
    } {
      owns_pet X _
    }
  }

  is_childless X {
    person X
    not parent X _
  }

  Adults = find P {
    person P
    is_adult P
  }

  dcg command (cmd Action Target) {
    verb Action
    noun Target
  }

  dcg verb (action find) { "find" }
  dcg noun (obj harry)   { "harry" }

  dcg sentence {
    noun_phrase
    verb_phrase
  }
  """
  |> Tpl.to_ast()
  |> TplProlog.to_prolog()

IO.puts("\n--- Generated Prolog ---")
IO.puts(prolog_code)
IO.puts("------------------------\n")