api.go

package cli

import (
	"errors"
	"fmt"
	"log/slog"

	"github.com/AlexsanderHamir/prof/engine/benchmark"
	"github.com/AlexsanderHamir/prof/engine/collector"
	"github.com/AlexsanderHamir/prof/engine/tracker"
	"github.com/AlexsanderHamir/prof/engine/version"
	"github.com/AlexsanderHamir/prof/internal/args"
	"github.com/AlexsanderHamir/prof/internal/config"
	"github.com/AlexsanderHamir/prof/internal/shared"
	"github.com/spf13/cobra"
)

var (
	// Root command flags.
	benchmarks []string
	profiles   []string
	tag        string
	count      int

	// Track command flags.
	baselineTag   string
	currentTag    string
	benchmarkName string
	profileType   string
	outputFormat  string
)

// CreateRootCmd creates and returns the root cobra command.
func CreateRootCmd() *cobra.Command {
	rootCmd := &cobra.Command{
		Use:   "prof",
		Short: "CLI tool for organizing pprof generated data, and analyzing performance differences at the profile level.",
		RunE:  runBenchmarks,
	}

	rootCmd.AddCommand(createProfManual())
	rootCmd.AddCommand(createProfAuto())
	rootCmd.AddCommand(createSetupCmd())
	rootCmd.AddCommand(createTrackCmd())
	rootCmd.AddCommand(createVersionCmd())

	return rootCmd
}

func createProfManual() *cobra.Command {
	manualCmd := &cobra.Command{
		Use:     shared.MANUALCMD,
		Short:   "Receives profile files and performs data collection and organization. (doesn't wrap go test)",
		Args:    cobra.MinimumNArgs(1),
		Example: fmt.Sprintf("prof %s --tag tagName cpu.prof memory.prof block.prof mutex.prof", shared.MANUALCMD),
		RunE: func(_ *cobra.Command, args []string) error {
			return collector.RunCollector(args, tag)
		},
	}

	tagFlag := "tag"
	manualCmd.Flags().StringVar(&tag, tagFlag, "", "The tag is used to organize the results")
	_ = manualCmd.MarkFlagRequired(tagFlag)

	return manualCmd
}

func createProfAuto() *cobra.Command {
	benchFlag := "benchmarks"
	profileFlag := "profiles"
	tagFlag := "tag"
	countFlag := "count"
	example := fmt.Sprintf(`prof %s --%s "BenchmarkGenPool" --%s "cpu,memory" --%s 10 --%s "tag1"`, shared.AUTOCMD, benchFlag, profileFlag, countFlag, tagFlag)

	cmd := &cobra.Command{
		Use:     shared.AUTOCMD,
		Short:   "Wraps `go test` and `pprof` to benchmark code and gather profiling data for performance investigations.",
		RunE:    runBenchmarks,
		Example: example,
	}

	cmd.Flags().StringSliceVar(&benchmarks, benchFlag, []string{}, `Benchmarks to run (e.g., "BenchmarkGenPool")"`)
	cmd.Flags().StringSliceVar(&profiles, profileFlag, []string{}, `Profiles to use (e.g., "cpu,memory,mutex")`)
	cmd.Flags().StringVar(&tag, tagFlag, "", "The tag is used to organize the results")
	cmd.Flags().IntVar(&count, countFlag, 0, "Number of runs")

	_ = cmd.MarkFlagRequired(benchFlag)
	_ = cmd.MarkFlagRequired(profileFlag)
	_ = cmd.MarkFlagRequired(tagFlag)
	_ = cmd.MarkFlagRequired(countFlag)

	return cmd
}

// createTrackCmd creates the track subcommand
func createTrackCmd() *cobra.Command {
	shortExplanation := "Compare performance between two benchmark runs to detect regressions and improvements"
	cmd := &cobra.Command{
		Use:   "track",
		Short: shortExplanation,
	}

	cmd.AddCommand(createTrackAutoCmd())
	cmd.AddCommand(createTrackManualCmd())

	return cmd
}

func createTrackAutoCmd() *cobra.Command {
	baseTagFlag := "base"
	currentTagFlag := "current"
	benchNameFlag := "bench-name"
	profileTypeFlag := "profile-type"
	outputFormatFlag := "output-format"
	example := fmt.Sprintf(`prof track auto --%s "tag1" --%s "tag2" --%s "cpu" --%s "BenchmarkGenPool" --%s "summary"`, baseTagFlag, currentTagFlag, profileTypeFlag, benchNameFlag, outputFormatFlag)
	longExplanation := fmt.Sprintf("This command only works if the %s command was used to collect and organize the benchmark and profile data, as it expects a specific directory structure generated by that process.", shared.AUTOCMD)
	shortExplanation := "If prof auto was used to collect the data, track auto can be used to analyze it, you just have to pass the tag name."

	cmd := &cobra.Command{
		Use:     shared.TrackAutoCMD,
		Short:   shortExplanation,
		Long:    longExplanation,
		RunE:    runTrackAuto,
		Example: example,
	}

	cmd.Flags().StringVar(&baselineTag, baseTagFlag, "", "Name of the baseline tag")
	cmd.Flags().StringVar(&currentTag, currentTagFlag, "", "Name of the current tag")
	cmd.Flags().StringVar(&benchmarkName, benchNameFlag, "", "Name of the benchmark")
	cmd.Flags().StringVar(&profileType, profileTypeFlag, "", "Profile type (cpu, memory, mutex, block)")
	cmd.Flags().StringVar(&outputFormat, outputFormatFlag, "detailed", `Output format: "summary" or "detailed"`)

	_ = cmd.MarkFlagRequired(baseTagFlag)
	_ = cmd.MarkFlagRequired(currentTagFlag)
	_ = cmd.MarkFlagRequired(benchNameFlag)
	_ = cmd.MarkFlagRequired(profileTypeFlag)

	return cmd
}

func createTrackManualCmd() *cobra.Command {
	baseFlag := "base"
	currentFlag := "current"
	outputFormatFlag := "output-format"
	example := fmt.Sprintf(`prof track %s --%s "path/to/profile_file.txt" --%s  "path/to/profile_file.txt"  --%s "summary"`, shared.TrackManualCMD, baseFlag, currentFlag, outputFormatFlag)

	cmd := &cobra.Command{
		Use:     shared.TrackManualCMD,
		Short:   "Manually specify the paths to the profile text files you want to compare.",
		RunE:    runTrackManual,
		Example: example,
	}

	cmd.Flags().StringVar(&baselineTag, baseFlag, "", "Name of the baseline tag")
	cmd.Flags().StringVar(&currentTag, currentFlag, "", "Name of the current tag")
	cmd.Flags().StringVar(&outputFormat, outputFormatFlag, "", "Output format choice choice")

	_ = cmd.MarkFlagRequired(baseFlag)
	_ = cmd.MarkFlagRequired(currentFlag)
	_ = cmd.MarkFlagRequired(outputFormatFlag)

	return cmd
}

func createVersionCmd() *cobra.Command {
	cmd := &cobra.Command{
		Use:                   "version",
		Short:                 "Shows the current version of prof and checks for updates.",
		RunE:                  runVersion,
		DisableFlagsInUseLine: true,
	}

	return cmd
}

// createSetupCmd creates the setup subcommand
func createSetupCmd() *cobra.Command {
	cmd := &cobra.Command{
		Use:                   "setup",
		Short:                 "Generates the template configuration file.",
		RunE:                  runSetup,
		DisableFlagsInUseLine: true,
	}

	return cmd
}

// Execute runs the CLI application
func Execute() error {
	return CreateRootCmd().Execute()
}

func runBenchmarks(_ *cobra.Command, _ []string) error {
	if len(benchmarks) == 0 {
		return errors.New("benchmarks flag is empty")
	}

	if len(profiles) == 0 {
		return errors.New("profiles flag is empty")
	}

	cfg, err := config.LoadFromFile(shared.ConfigFilename)
	if err != nil {
		cfg = &config.Config{}
	}

	if err = benchmark.SetupDirectories(tag, benchmarks, profiles); err != nil {
		return fmt.Errorf("failed to setup directories: %w", err)
	}

	benchArgs := &args.BenchArgs{
		Benchmarks: benchmarks,
		Profiles:   profiles,
		Count:      count,
		Tag:        tag,
	}

	printConfiguration(benchArgs, cfg.FunctionFilter)

	if err = runBenchAndGetProfiles(benchArgs, cfg.FunctionFilter); err != nil {
		return err
	}

	return nil
}

// runVersion handles the version command execution
func runVersion(_ *cobra.Command, _ []string) error {
	current, latest := version.Check()
	output := version.FormatOutput(current, latest)
	fmt.Print(output)
	return nil
}

// runSetup handles the setup command execution
func runSetup(_ *cobra.Command, _ []string) error {
	return config.CreateTemplate()
}

var validFormats = map[string]bool{
	"summary":       true,
	"detailed":      true,
	"summary-html":  true,
	"detailed-html": true,
}

// runTrack handles the track command execution
func runTrackAuto(_ *cobra.Command, _ []string) error {
	if !validFormats[outputFormat] {
		return fmt.Errorf("invalid output format '%s'. Valid formats: summary, detailed", outputFormat)
	}

	report, err := tracker.CheckPerformanceDifferences(baselineTag, currentTag, benchmarkName, profileType)
	if err != nil {
		return fmt.Errorf("failed to track performance differences: %w", err)
	}

	noFunctionChanges := len(report.FunctionChanges) == 0
	if noFunctionChanges {
		slog.Info("No function changes detected between the two runs")
		return nil
	}

	chooseOutputFormat(report)

	return nil
}

func runTrackManual(_ *cobra.Command, _ []string) error {
	if !validFormats[outputFormat] {
		return fmt.Errorf("invalid output format '%s'. Valid formats: summary, detailed", outputFormat)
	}

	report, err := tracker.CheckPerformanceDifferencesManual(baselineTag, currentTag)
	if err != nil {
		return fmt.Errorf("failed to track performance differences: %w", err)
	}

	noFunctionChanges := len(report.FunctionChanges) == 0
	if noFunctionChanges {
		slog.Info("No function changes detected between the two runs")
		return nil
	}

	chooseOutputFormat(report)

	return nil
}