defmodule Mix.Tasks.Vibe.Docs do
  use Mix.Task
  alias Vibe.DocsHelper

  @shortdoc "Get documentation for modules and functions"
  @moduledoc """
  Fetches and displays documentation for modules and functions.

  ## Usage

      mix vibe.docs [--list-modules]
      mix vibe.docs MODULE_NAME
      mix vibe.docs MODULE_NAME.function_name[/arity]

  ## Examples

      mix vibe.docs --list-modules
      mix vibe.docs Enum
      mix vibe.docs Enum.map
      mix vibe.docs Enum.map/2

  ## Options

      --format FORMAT     Output format (text, json, markdown, default: config value or markdown)
      --list-modules      List all available modules
      --help, -h          Show this help message

  """

  @impl Mix.Task
  def run(args) do
    {opts, remaining_args, _} =
      OptionParser.parse(args,
        strict: [format: :string, list_modules: :boolean, help: :boolean],
        aliases: [h: :help]
      )

    # Make sure the project is compiled
    Mix.Task.run("compile")

    format = opts_to_format(opts)

    cond do
      Keyword.has_key?(opts, :help) ->
        print_help()

      Keyword.has_key?(opts, :list_modules) ->
        list_modules(format)

      length(remaining_args) > 0 ->
        module_and_function = List.first(remaining_args)
        # Parse module and function
        case parse_module_and_function(module_and_function) do
          {:module, module} ->
            get_module_docs(module, format)

          {:function, module, function, arity} ->
            get_function_docs(module, function, arity, format)
        end

      true ->
        print_help()
    end
  end

  defp opts_to_format(opts) do
    cond do
      opts[:format] -> String.to_atom(opts[:format])
      true -> Vibe.output_format()
    end
  end

  defp list_modules(format) do
    modules = DocsHelper.list_project_modules()

    output =
      case format do
        :json ->
          Jason.encode!(%{modules: Enum.map(modules, &Atom.to_string/1)}, pretty: true)

        :markdown ->
          module_list =
            modules
            |> Enum.sort()
            |> Enum.map(&"- `#{inspect(&1)}`")
            |> Enum.join("\n")

          """
          # Available Modules

          #{module_list}
          """

        _ ->
          modules
          |> Enum.sort()
          |> Enum.map(&inspect/1)
          |> Enum.join("\n")
      end

    IO.puts(output)
  end

  defp get_module_docs(module, format) do
    case DocsHelper.get_module_docs(module) do
      {:ok, docs} ->
        output =
          case format do
            :json ->
              Jason.encode!(docs, pretty: true)

            :markdown ->
              functions_list =
                docs.functions
                |> Map.values()
                |> Enum.sort_by(& &1.name)
                |> Enum.map(fn f -> "- `#{f.name}/#{f.arity}` - #{extract_first_line(f.doc)}" end)
                |> Enum.join("\n")

              """
              # Module `#{inspect(docs.module)}`

              #{docs.module_doc || "*No module documentation available*"}

              ## Functions

              #{functions_list}
              """

            _ ->
              module_str = "MODULE: #{inspect(docs.module)}\n"
              doc_str = if docs.module_doc, do: "#{docs.module_doc}\n\n", else: "\n"

              functions_str =
                docs.functions
                |> Map.values()
                |> Enum.sort_by(& &1.name)
                |> Enum.map(fn f -> "  #{f.name}/#{f.arity}" end)
                |> Enum.join("\n")

              "#{module_str}#{doc_str}FUNCTIONS:\n#{functions_str}"
          end

        IO.puts(output)

      {:error, message} ->
        print_error(message, format)
    end
  end

  defp get_function_docs(module, function, arity, format) do
    case DocsHelper.get_function_docs(module, function, arity) do
      {:ok, docs} ->
        output =
          case format do
            :json ->
              Jason.encode!(docs, pretty: true)

            :markdown ->
              docs
              |> Enum.map(fn doc ->
                """
                # `#{inspect(doc.module)}.#{doc.name}/#{doc.arity}`

                ```elixir
                #{doc.name}(#{function_args(doc.arity)})
                ```

                #{doc.doc || "*No documentation available*"}
                """
              end)
              |> Enum.join("\n\n---\n\n")

            _ ->
              docs
              |> Enum.map(fn doc ->
                """
                #{inspect(doc.module)}.#{doc.name}/#{doc.arity}
                #{doc.doc || "No documentation available"}
                """
              end)
              |> Enum.join("\n\n")
          end

        IO.puts(output)

      {:error, message} ->
        print_error(message, format)
    end
  end

  defp print_error(message, format) do
    output =
      case format do
        :json -> Jason.encode!(%{error: message}, pretty: true)
        :markdown -> "**Error:** #{message}"
        _ -> "Error: #{message}"
      end

    IO.puts(output)
  end

  defp function_args(0), do: ""

  defp function_args(arity) do
    1..arity
    |> Enum.map(fn i -> "arg#{i}" end)
    |> Enum.join(", ")
  end

  defp extract_first_line(doc) when is_binary(doc) do
    doc
    |> String.split("\n", parts: 2)
    |> List.first()
    |> String.trim()
  end

  defp extract_first_line(_), do: "*No documentation available*"

  defp parse_module_and_function(input) do
    cond do
      # Check for Module.function/arity format
      String.match?(input, ~r/^[A-Za-z0-9_.]+\.[a-zA-Z0-9_!?]+\/[0-9]+$/) ->
        [module_and_function, arity] = String.split(input, "/")
        [module, function] = split_module_and_function(module_and_function)
        {:function, module, function, String.to_integer(arity)}

      # Check for Module.function format
      String.match?(input, ~r/^[A-Za-z0-9_.]+\.[a-zA-Z0-9_!?]+$/) ->
        [module, function] = split_module_and_function(input)
        {:function, module, function, nil}

      # Assume it's just a module
      true ->
        {:module, input}
    end
  end

  defp split_module_and_function(input) do
    parts = String.split(input, ".")
    function = List.last(parts)
    module = Enum.join(Enum.drop(parts, -1), ".")
    [module, function]
  end

  defp print_help do
    IO.puts("mix vibe.docs - Get documentation for modules and functions")
    IO.puts("")
    IO.puts("Usage:")
    IO.puts("  mix vibe.docs [--list-modules]")
    IO.puts("  mix vibe.docs MODULE_NAME")
    IO.puts("  mix vibe.docs MODULE_NAME.function_name[/arity]")
    IO.puts("")
    IO.puts("Examples:")
    IO.puts("  mix vibe.docs --list-modules")
    IO.puts("  mix vibe.docs Enum")
    IO.puts("  mix vibe.docs Enum.map")
    IO.puts("  mix vibe.docs Enum.map/2")
    IO.puts("")
    IO.puts("Options:")
    IO.puts("  --format FORMAT     Output format (text, json, markdown, default: config value or markdown)")
    IO.puts("  --list-modules      List all available modules")
    IO.puts("  --help, -h          Show this help message")
    IO.puts("")
    IO.puts("Feedback:")
    IO.puts("  📚 Groovy docs for your coding journey! 🎵")
  end
end