defmodule Mix.Tasks.Compile.Protos do
  @moduledoc """
  Mix task to compile protobuf files into Elixir modules with gRPC support.
  
  This task uses the protoc compiler to generate Elixir code from .proto files.
  It includes support for gRPC service definitions.
  """
  use Mix.Task

  @shortdoc "Compiles protobuf files to Elixir code with gRPC support"

  @impl Mix.Task
  def run(_args) do
    # Ensure protoc is installed
    if !protoc_installed?() do
      Mix.raise("protoc command not found. Please install Protocol Buffers compiler.")
    end

    # Check if protoc-gen-elixir is installed
    if !protoc_gen_elixir_installed?() do
      Mix.shell().info("Installing protoc-gen-elixir plugin...")
      install_protoc_gen_elixir()
    end

    # Define paths
    project_dir = Mix.Project.deps_path() |> Path.dirname()
    proto_dir = Path.join(project_dir, "priv/proto")
    output_dir = Path.join(project_dir, "lib/generated")
    
    # Create output directory if it doesn't exist
    File.mkdir_p!(output_dir)

    # Download and extract Google API proto files if needed
    google_api_dir = Path.join(proto_dir, "google/api")
    if !File.exists?(google_api_dir) do
      download_google_api_protos(proto_dir)
    end

    # Compile the main proto files needed for BES implementation, NativeLink service, and Bazel Remote Execution API
    proto_files = [
      "build_events.proto",
      "build_event_stream.proto",
      "publish_build_event.proto",
      "nativelink.proto",
      "build/bazel/remote/execution/v2/remote_execution.proto",
      "build/bazel/semver/semver.proto",
      "google/bytestream/bytestream.proto",
      "google/longrunning/operations.proto",
      "google/rpc/status.proto"
    ]
    
    compile_proto_files(proto_dir, output_dir, proto_files)
    
    :ok
  end
  
  defp compile_proto_files(proto_dir, output_dir, files) do
    # Set up proto paths for correct import resolution
    # 1. Main proto dir
    # 2. The parent directory of proto dir to handle imports like "proto/file.proto"
    parent_proto_dir = Path.join(proto_dir, "..")
    proto_path_args = [
      "--proto_path=#{proto_dir}",
      "--proto_path=#{parent_proto_dir}"
    ]
    
    # For imports like "google/api/annotations.proto", add include paths that might contain these
    include_proto_dirs = [
      "/usr/local/include",  # Common location for well-known protos
      "/usr/include"       # Alternative location
    ]
    
    proto_path_args = proto_path_args ++ Enum.map(include_proto_dirs, fn dir -> "--proto_path=#{dir}" end)
    
    # Process each proto file
    Enum.each(files, fn file ->
      file_path = Path.join(proto_dir, file)
      
      if File.exists?(file_path) do
        proto_paths = Enum.join(proto_path_args, " ")
        cmd = "protoc --elixir_out=plugins=grpc:#{output_dir} #{proto_paths} #{file_path}"
        
        Mix.shell().info("Compiling #{file}...")
        Mix.shell().info("Command: #{cmd}")
        Mix.shell().info("File exists check: #{File.exists?(file_path)}")
        
        case System.cmd("sh", ["-c", cmd], stderr_to_stdout: true) do
          {output, 0} ->
            Mix.shell().info("Successfully compiled #{file}")
            if String.trim(output) != "", do: Mix.shell().info(output)
            # Check if the output file was generated
            expected_module_name = file
              |> Path.basename(".proto")
              |> Macro.camelize()
            Mix.shell().info("Generated output path: #{output_dir}")
            Mix.shell().info("Generated files check: #{inspect(File.ls!(output_dir))}")
          {error, code} ->
            Mix.shell().error("Failed to compile #{file}: error code #{code}")
            Mix.shell().error(error)
        end
      else
        Mix.shell().error("Proto file not found: #{file_path}")
      end
    end)
  end

  defp download_google_api_protos(proto_dir) do
    Mix.shell().info("Setting up Google API proto files...")
    
    # Create temporary directory for download
    tmp_dir = Path.join(System.tmp_dir!(), "google_api_protos")
    File.mkdir_p!(tmp_dir)
    
    # Clone the googleapis repository to get the required protos
    cmd = "git clone --depth 1 https://github.com/googleapis/googleapis.git #{tmp_dir}"
    Mix.shell().info("Downloading Google API protos: #{cmd}")
    
    case System.cmd("sh", ["-c", cmd], stderr_to_stdout: true) do
      {output, 0} ->
        # Create the target directories
        google_dir = Path.join(proto_dir, "google")
        File.mkdir_p!(Path.join(google_dir, "api"))
        
        # Copy the required google/api/*.proto files
        api_protos = Path.wildcard(Path.join([tmp_dir, "google/api/*.proto"]))
        Enum.each(api_protos, fn proto_file ->
          filename = Path.basename(proto_file)
          dest_file = Path.join([google_dir, "api", filename])
          File.cp!(proto_file, dest_file)
          Mix.shell().info("Copied google/api/#{filename}")
        end)

        # Also copy google/bytestream, google/longrunning, and google/rpc directories
        additional_dirs = ["bytestream", "longrunning", "rpc", "protobuf"]
        Enum.each(additional_dirs, fn dir_name ->
          src_dir = Path.join([tmp_dir, "google", dir_name])
          dest_dir = Path.join([google_dir, dir_name])
          if File.dir?(src_dir) do
            File.mkdir_p!(dest_dir)
            File.cp_r!(src_dir, dest_dir)
            Mix.shell().info("Copied google/#{dir_name} protos")
          else
            Mix.shell().error("Directory #{src_dir} not found in googleapis repo")
          end
        end)
        
        # Clean up
        File.rm_rf!(tmp_dir)
        Mix.shell().info("Google API protos setup complete")
        
      {error, _code} ->
        Mix.shell().error("Failed to download Google API protos: #{error}")
        Mix.shell().error("You may need to manually install Google API proto files")
    end
  end
  
  defp protoc_installed? do
    case System.cmd("which", ["protoc"], stderr_to_stdout: true) do
      {_, 0} -> true
      _ -> false
    end
  end
  
  defp protoc_gen_elixir_installed? do
    case System.cmd("which", ["protoc-gen-elixir"], stderr_to_stdout: true) do
      {_, 0} -> true
      _ -> false
    end
  end
  
  defp install_protoc_gen_elixir do
    Mix.shell().info("Installing protoc-gen-elixir via Mix...")
    System.cmd("mix", ["escript.install", "hex", "protobuf", "--force"], stderr_to_stdout: true)
    Mix.shell().info("protoc-gen-elixir installed successfully")
  end
end
