Code.require_file("tpl.exs")

defmodule TplLua do
  def transpile(ast_nodes) do
    ast_nodes
    |> Enum.map(&transpile_node(&1, 0))
    |> Enum.reject(&(&1 == ""))
    |> Enum.join("\n\n")
  end

  # Filters out noise tokens that should not appear as arguments or statements
  defp relevant?(nil), do: false
  defp relevant?(%{t: t}) when t in [:newline, :comment], do: false
  defp relevant?(_), do: true

  # Coalesces `fn arg... brace` sequences in a flat token list into fn sexp nodes.
  # Needed when `fn` appears as an argument inside a larger call, e.g. `map "k" fn { ... }`.
  defp coalesce_fns([]), do: []

  defp coalesce_fns([%{t: :token, str: "fn"} | rest]) do
    {fn_args, remainder} = Enum.split_while(rest, fn n -> n[:t] not in [:brace] end)
    case remainder do
      [brace | after_brace] ->
        fn_node = %{t: :sexp, list: [%{t: :token, str: "fn"} | fn_args ++ [brace]]}
        [fn_node | coalesce_fns(after_brace)]
      [] ->
        [%{t: :token, str: "fn"} | coalesce_fns(rest)]
    end
  end

  defp coalesce_fns([head | rest]), do: [head | coalesce_fns(rest)]

  defp transpile_node(node, indent) do
    indent_str = String.duplicate("  ", indent)

    case node do
      # Bug 9: comment nodes → Lua -- comments
      %{t: :comment, str: str} ->
        str
        |> String.split("\n")
        |> Enum.map(&"#{indent_str}-- #{String.trim(&1)}")
        |> Enum.join("\n")

      # Bug 2 & 6: nil / noise / empty brace
      nil -> ""
      %{t: t} when t in [:newline, :dot] -> ""

      # Bug 6: empty unnamed brace → empty Lua table
      %{t: :brace, list: []} -> "{}"

      %{t: :sexp, list: list} ->
        transpile_sexp(list, indent)

      %{t: :brace, list: list, name: %{str: "t"}} ->
        transpile_table(list, indent)

      %{t: :brace, list: list, name: %{str: "l"}} ->
        transpile_list(list, indent)

      %{t: :brace, list: list} ->
        transpile_block(list, indent, :block)

      %{t: :token, str: str} ->
        str

      %{t: :string, chunks: chunks, values: values} ->
        transpile_string(chunks, values, indent)

      other ->
        " --[[ UNHANDLED NODE: #{inspect(other, pretty: true)} ]]--"
    end
  end

  # Bug 7: non-brace fn body gets `return` + inner indentation
  defp transpile_node(node, indent, context) do
    case node do
      %{t: :brace, list: list} ->
        transpile_block(list, indent, context)

      _ ->
        body_str = transpile_node(node, indent)

        case context do
          :function_body ->
            inner = String.duplicate("  ", indent + 1)
            "#{inner}return #{body_str}"

          _ ->
            body_str
        end
    end
  end

  defp transpile_sexp(list, indent) do
    indent_str = String.duplicate("  ", indent)

    case list do
      # Assignment: = lhs rhs
      [%{str: "="}, %{t: :sexp, list: [%{str: "l"}, %{str: name}]}, rhs] ->
        "local #{name} = #{transpile_node(rhs, indent)}"

      [%{str: "="}, lhs, rhs] ->
        "#{transpile_node(lhs, indent)} = #{transpile_node(rhs, indent)}"

      # Bug 4: if form → if/then/end
      [%{str: "if"} | rest] ->
        filtered = Enum.filter(rest, &relevant?/1)
        {cond_parts, [body]} = Enum.split(filtered, -1)
        cond_str = cond_parts |> Enum.map(&transpile_node(&1, indent)) |> Enum.join(" ")
        body_str = transpile_node(body, indent, :block)
        "if #{cond_str} then\n#{body_str}\n#{indent_str}end"

      # Bug 3: a : .method args → a.method(args)
      [%{str: ":"}, lhs, %{t: :sexp, list: [%{str: "."}, nil, rhs]}] ->
        lhs_str = transpile_node(lhs, indent)

        case rhs do
          %{t: :sexp, list: [method | args]} ->
            method_str = transpile_node(method, indent)
            arg_list = args |> Enum.filter(&relevant?/1) |> Enum.map(&transpile_node(&1, indent)) |> Enum.join(", ")
            "#{lhs_str}.#{method_str}(#{arg_list})"

          prop ->
            "#{lhs_str}.#{transpile_node(prop, indent)}"
        end

      # Pipe: a : b c → a:b(c)
      [%{str: ":"}, lhs, %{t: :sexp, list: [func | args]}] ->
        lhs_str = transpile_node(lhs, indent)
        func_str = transpile_node(func, indent)
        # Bug 2: filter noise from args
        arg_list = args |> Enum.filter(&relevant?/1) |> Enum.map(&transpile_node(&1, indent)) |> Enum.join(", ")
        "#{lhs_str}:#{func_str}(#{arg_list})"

      # Pipe: a : b → a:b()
      [%{str: ":"}, lhs, rhs] ->
        "#{transpile_node(lhs, indent)}:#{transpile_node(rhs, indent)}()"

      # Dot access: . lhs rhs
      [%{str: "."}, lhs, rhs] ->
        "#{transpile_node(lhs, indent)}.#{transpile_node(rhs, indent)}"

      # Anonymous function: fn args... body
      [%{str: "fn"} | rest] ->
        filtered = Enum.filter(rest, &relevant?/1)
        {args, [body]} = Enum.split(filtered, -1)
        arg_list = args |> Enum.map(&transpile_node(&1, indent)) |> Enum.join(", ")
        body_str = transpile_node(body, indent, :function_body)
        "function(#{arg_list})\n#{body_str}\n#{indent_str}end"

      # For loop: for i v in (ipairs arr) { body }
      [%{str: "for"} | rest] ->
        filtered = Enum.filter(rest, &relevant?/1)
        {vars_and_in, [body_node]} = Enum.split(filtered, -1)
        {vars, [%{str: "in"}, call_node]} = Enum.split(vars_and_in, -2)
        var_list = vars |> Enum.map(&transpile_node(&1, indent)) |> Enum.join(", ")
        body_str = transpile_node(body_node, indent, :block)
        "for #{var_list} in #{transpile_node(call_node, indent)} do\n#{body_str}\n#{indent_str}end"

      # Function call: func arg1 arg2 ...
      [func | args] ->
        func_str = transpile_node(func, indent)
        arg_list =
          args
          |> Enum.filter(&relevant?/1)
          |> coalesce_fns()
          |> Enum.map(&transpile_node(&1, indent))
          |> Enum.join(", ")
        "#{func_str}(#{arg_list})"

      [] ->
        "()"
    end
  end

  # True when a node is a block-level statement (if/for/while) that cannot be returned
  defp block_stmt?(%{t: :sexp, list: [%{str: kw} | _]}) when kw in ["if", "for", "while"],
    do: true

  defp block_stmt?(_), do: false

  defp transpile_block(list, indent, context) do
    inner = String.duplicate("  ", indent + 1)

    statements =
      list
      |> Enum.map(&transpile_node(&1, indent + 1))
      |> Enum.reject(&(&1 == ""))
      |> Enum.map(&"#{inner}#{&1}")

    case context do
      # Bug 8: don't prepend `return` to if/for/while
      :function_body ->
        last_node = list |> Enum.filter(&relevant?/1) |> List.last()

        if last_node == nil or block_stmt?(last_node) do
          statements |> Enum.join("\n")
        else
          case statements do
            [] ->
              ""

            _ ->
              {body, [last]} = Enum.split(statements, -1)
              return_last = String.replace(last, ~r/^\s*/, "#{inner}return ", global: false)
              (body ++ [return_last]) |> Enum.join("\n")
          end
        end

      _ ->
        statements |> Enum.join("\n")
    end
  end

  defp transpile_list(list, indent) do
    if list == [] do
      "{}"
    else
      inner = String.duplicate("  ", indent + 1)
      elements = list |> Enum.map(&"#{inner}#{transpile_node(&1, indent + 1)}") |> Enum.join(",\n")
      "{\n#{elements}\n#{String.duplicate("  ", indent)}}"
    end
  end

  defp transpile_table(list, indent) do
    if list == [] do
      "{}"
    else
      inner = String.duplicate("  ", indent + 1)

      pairs =
        list
        |> Enum.flat_map(fn
          %{t: :sexp, list: kv_list} ->
            kv_list
            |> Enum.filter(&relevant?/1)
            |> Enum.chunk_every(2)
            |> Enum.map(fn
              [k, v] -> "#{inner}#{transpile_node(k, indent + 1)} = #{transpile_node(v, indent + 1)}"
              [k] -> "#{inner}#{transpile_node(k, indent + 1)} = true"
            end)

          other ->
            ["#{inner}--[[ INVALID TABLE ENTRY: #{inspect(other)} ]]--"]
        end)
        |> Enum.join(",\n")

      "{\n#{pairs}\n#{String.duplicate("  ", indent)}}"
    end
  end

  # Bug 5: unwrap brace so interpolated values aren't treated as indented blocks
  defp transpile_interp_value(%{t: :brace, list: [single]}, indent),
    do: transpile_node(single, indent)

  defp transpile_interp_value(%{t: :brace, list: list}, indent) when list != [],
    do: transpile_node(%{t: :sexp, list: list}, indent)

  defp transpile_interp_value(other, indent),
    do: transpile_node(other, indent)

  defp transpile_string(chunks, values, indent) do
    if values == [] do
      "\"#{Enum.at(chunks, 0, "")}\""
    else
      parts =
        Enum.zip(chunks, values)
        |> Enum.flat_map(fn {chunk, value} ->
          ["\"#{chunk}\"", "tostring(#{transpile_interp_value(value, indent)})"]
        end)

      final_parts =
        if length(chunks) > length(values),
          do: parts ++ ["\"#{List.last(chunks)}\""],
          else: parts

      final_parts
      |> Enum.reject(&(&1 == "\"\""))
      |> Enum.join(" .. ")
    end
  end
end

defmodule TplLua.CLI do
  @moduledoc """
  Provides a command-line interface for transpiling `.lua.tpl` files.
  """

  @doc """
  Transpiles `.lua.tpl` source files or directories into `.lua` files.

  This function can be used in two ways:

  1.  **Directory to Directory**:
      `run(input_dir, output_dir)`
      Recursively finds all `*.lua.tpl` files in `input_dir`, transpiles them,
      and writes the output `.lua` files to `output_dir`, preserving the
      directory structure.

  2.  **File to File**:
      `run(input_file, output_file)`
      Transpiles a single `input_file` and writes the result to `output_file`.

  Returns `:ok` on success, or `{:error, reason}` on failure.
  """
  def run(input_path, output_path) do
    cond do
      File.dir?(input_path) ->
        transpile_directory(input_path, output_path)

      File.regular?(input_path) ->
        transpile_single_file(input_path, output_path)

      true ->
        {:error, "Input path not found: #{input_path}"}
    end
  end

  defp transpile_directory(input_dir, output_dir) do
    input_dir = Path.expand(input_dir)
    output_dir = Path.expand(output_dir)

    File.mkdir_p!(output_dir)

    input_dir
    |> Path.join("**/*.lua.tpl")
    |> Path.wildcard()
    |> Enum.each(fn source_file ->
      relative_path = String.replace(source_file, input_dir <> "/", "")
      dest_file = Path.join(output_dir, relative_path) |> String.replace(".lua.tpl", ".lua")
      File.mkdir_p!(Path.dirname(dest_file))
      transpile_file(source_file, dest_file)
    end)

    :ok
  end

  defp transpile_single_file(input_file, output_path) do
    dest_file =
      if File.dir?(output_path) do
        # If output is a directory, create file with same name inside it
        Path.join(output_path, Path.basename(input_file)) |> String.replace(".lua.tpl", ".lua")
      else
        # Otherwise, use the exact output path
        output_path
      end

    File.mkdir_p!(Path.dirname(dest_file))
    transpile_file(input_file, dest_file)
    :ok
  end

  defp transpile_file(source, dest) do
    IO.puts("Transpiling #{source} -> #{dest}")

    source
    |> File.read!()
    |> Tpl.to_ast()
    |> TplLua.transpile()
    |> then(&File.write!(dest, &1))
  end
end

example =
  """
   ---

  on_ff = fn a b c (vim.lsp.buf.format)
  on_ff = fn a b c { 
    (vim.lsp.buf.format)
  } 
   ---

  map "n" "<leader>ff" (fn a b c (vim.lsp.buf.format)) t{ desc "Format code" }

   ---

  l asd = fn x y z {
  print "Hello, World!"
  }
   ---

  player.add_hook = fn self event data {
  self:process_event event data
  }
   ---

  arr_map = fn arr fun {
  l result = {}
  for i v in (ipairs arr) {
    table.insert result (fun v)
  }
  result
  }
  ---
  day = fetch_day db
  name = fetch_name user_id
  greeting = "Hello, {name}! Today is {day}."
  ---
  asd = fn {print "Hello"}
  ---
  asd = t{ key1 "value1" key2 "value2" }
  ---

  map "n" "<leader>stf" fn {
    l timer_minutes = (choose_timer)
    if not timer_minutes {
      print "No timer selected"
    }
    require "grug-far"
      : .open t{
        transient true
        prefills t{
          timers (timer_minutes : at 1 : tostring)
        }
      }
  }
  """

# Keep the interactive example runner for development
example
|> String.split("---\n")
|> Enum.each(fn input ->
  IO.puts("===================== Input ====")
  IO.puts(input)
  IO.puts("==== Parsed ====")

  state =
    input
    |> Tpl.chunk()

  # |> IO.inspect()
  # |> print_chunkd()

  chunks = Map.get(state, :program_output)
  IO.puts("==== AST ====")

  lua_code =
    chunks
    |> Tpl.chunks_to_ast()
    |> TplLua.transpile()

  IO.puts("\n==== Transpiled Lua ====")

  lua_code
  |> IO.puts()

  IO.puts("\n\n")
end)

# --- Example usage for the new CLI ---
#
# To run this, you could create some directories and files:
#
# mkdir -p lua_tpl/plugins
# echo 'l my_plugin = fn { print "plugin loaded" }' > lua_tpl/plugins/my_plugin.lua.tpl
# echo 'l main = fn { print "init file" }' > lua_tpl/init.lua.tpl
#
# Then, you could uncomment and run the following line in iex or from the script:
#
# TplLua.CLI.run("lua_tpl", "lua_output")
#
# This would produce:
# lua_output/plugins/my_plugin.lua
# lua_output/init.lua
TplLua.CLI.run("lua_tpl", "lua_output")