helpers.go

package cli

import (
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/AlecAivazis/survey/v2"
	"github.com/AlexsanderHamir/prof/engine/benchmark"
	"github.com/AlexsanderHamir/prof/engine/collector"

	"github.com/AlexsanderHamir/prof/engine/tracker"
	"github.com/AlexsanderHamir/prof/internal"
	"github.com/spf13/cobra"
)

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

	// Track command flags.
	Baseline            string
	Current             string
	benchmarkName       string
	profileType         string
	outputFormat        string
	failOnRegression    bool
	regressionThreshold float64
)

const (
	tuiPageSize          = 20
	minTagsForComparison = 2
)

// 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.",
	}

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

	return rootCmd
}

func createProfManual() *cobra.Command {
	manualCmd := &cobra.Command{
		Use:     internal.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", internal.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"`, internal.AUTOCMD, benchFlag, profileFlag, countFlag, tagFlag)

	cmd := &cobra.Command{
		Use:   internal.AUTOCMD,
		Short: "Wraps `go test` and `pprof` to benchmark code and gather profiling data for performance investigations.",
		RunE: func(_ *cobra.Command, _ []string) error {
			return benchmark.RunBenchmarks(benchmarks, profiles, tag, count)
		},
		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
}

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"
	currentFlag := "current"
	benchNameFlag := "bench-name"
	profileTypeFlag := "profile-type"
	outputFormatFlag := "output-format"
	failFlag := "fail-on-regression"
	thresholdFlag := "regression-threshold"
	example := fmt.Sprintf(`prof track auto --%s "tag1" --%s "tag2" --%s "cpu" --%s "BenchmarkGenPool" --%s "summary"`, baseTagFlag, currentFlag, 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.", internal.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:   internal.TrackAutoCMD,
		Short: shortExplanation,
		Long:  longExplanation,
		RunE: func(_ *cobra.Command, _ []string) error {
			selections := &tracker.Selections{
				OutputFormat:        outputFormat,
				Baseline:            Baseline,
				Current:             Current,
				ProfileType:         profileType,
				BenchmarkName:       benchmarkName,
				RegressionThreshold: regressionThreshold,
				UseThreshold:        failOnRegression,
			}
			return tracker.RunTrackAuto(selections)
		},
		Example: example,
	}

	cmd.Flags().StringVar(&Baseline, baseTagFlag, "", "Name of the baseline tag")
	cmd.Flags().StringVar(&Current, currentFlag, "", "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.Flags().BoolVar(&failOnRegression, failFlag, false, "Exit with non-zero code if regression exceeds threshold")
	cmd.Flags().Float64Var(&regressionThreshold, thresholdFlag, 0.0, "Fail when worst flat regression exceeds this percent (e.g., 5.0)")

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

	return cmd
}

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

	cmd := &cobra.Command{
		Use:   internal.TrackManualCMD,
		Short: "Manually specify the paths to the profile text files you want to compare.",
		RunE: func(_ *cobra.Command, _ []string) error {
			selections := &tracker.Selections{
				OutputFormat:        outputFormat,
				Baseline:            Baseline,
				Current:             Current,
				ProfileType:         profileType,
				BenchmarkName:       benchmarkName,
				RegressionThreshold: regressionThreshold,
				UseThreshold:        failOnRegression,
				IsManual:            true,
			}
			return tracker.RunTrackManual(selections)
		},
		Example: example,
	}

	cmd.Flags().StringVar(&Baseline, baseFlag, "", "Name of the baseline tag")
	cmd.Flags().StringVar(&Current, currentFlag, "", "Name of the current tag")
	cmd.Flags().StringVar(&outputFormat, outputFormatFlag, "", "Output format choice choice")
	cmd.Flags().BoolVar(&failOnRegression, failFlag, false, "Exit with non-zero code if regression exceeds threshold")
	cmd.Flags().Float64Var(&regressionThreshold, thresholdFlag, 0.0, "Fail when worst flat regression exceeds this percent (e.g., 5.0)")

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

	return cmd
}

func createSetupCmd() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "setup",
		Short: "Generates the template configuration file.",
		RunE: func(_ *cobra.Command, _ []string) error {
			return internal.CreateTemplate()
		},
		DisableFlagsInUseLine: true,
	}

	return cmd
}

// getTrackSelections collects all user selections interactively
func getTrackSelections(tags []string) (*tracker.Selections, error) {
	selections := &tracker.Selections{}

	// Select baseline tag
	baselinePrompt := &survey.Select{
		Message:  "Select baseline tag (the 'before' version) [Press Enter to select]:",
		Options:  tags,
		PageSize: tuiPageSize,
	}
	if err := survey.AskOne(baselinePrompt, &selections.Baseline, survey.WithValidator(survey.Required)); err != nil {
		return nil, err
	}

	// Select current tag (filter out baseline)
	var currentOptions []string
	for _, tag := range tags {
		if tag != selections.Baseline {
			currentOptions = append(currentOptions, tag)
		}
	}

	currentPrompt := &survey.Select{
		Message:  "Select current tag (the 'after' version) [Press Enter to select]:",
		Options:  currentOptions,
		PageSize: tuiPageSize,
	}
	if err := survey.AskOne(currentPrompt, &selections.Current, survey.WithValidator(survey.Required)); err != nil {
		return nil, err
	}

	// Discover and select benchmark
	if err := selectBenchmark(selections); err != nil {
		return nil, err
	}

	// Discover and select profile type
	if err := selectProfileType(selections); err != nil {
		return nil, err
	}

	// Select output format
	if err := selectOutputFormat(selections); err != nil {
		return nil, err
	}

	// Ask about regression threshold
	if err := selectRegressionThreshold(selections); err != nil {
		return nil, err
	}

	return selections, nil
}

// selectBenchmark discovers and selects a benchmark
func selectBenchmark(selections *tracker.Selections) error {
	availableBenchmarks, err := discoverAvailableBenchmarks(selections.Baseline)
	if err != nil {
		return fmt.Errorf("failed to discover benchmarks for tag %s: %w", selections.Baseline, err)
	}
	if len(availableBenchmarks) == 0 {
		return fmt.Errorf("no benchmarks found for tag %s", selections.Baseline)
	}

	benchPrompt := &survey.Select{
		Message:  "Select benchmark to compare [Press Enter to select]:",
		Options:  availableBenchmarks,
		PageSize: tuiPageSize,
	}
	return survey.AskOne(benchPrompt, &selections.BenchmarkName, survey.WithValidator(survey.Required))
}

// selectProfileType discovers and selects a profile type
func selectProfileType(selections *tracker.Selections) error {
	availableProfiles, err := discoverAvailableProfiles(selections.Baseline, selections.BenchmarkName)
	if err != nil {
		return fmt.Errorf("failed to discover profiles for tag %s, benchmark %s: %w", selections.Baseline, selections.BenchmarkName, err)
	}
	if len(availableProfiles) == 0 {
		return fmt.Errorf("no profiles found for tag %s, benchmark %s", selections.Baseline, selections.BenchmarkName)
	}

	profilePrompt := &survey.Select{
		Message:  "Select profile type to compare [Press Enter to select]:",
		Options:  availableProfiles,
		PageSize: tuiPageSize,
	}
	return survey.AskOne(profilePrompt, &selections.ProfileType, survey.WithValidator(survey.Required))
}

// selectOutputFormat selects the output format
func selectOutputFormat(selections *tracker.Selections) error {
	outputFormats := []string{"summary", "detailed", "summary-html", "detailed-html", "summary-json", "detailed-json"}
	formatPrompt := &survey.Select{
		Message:  "Select output format [Press Enter to select]:",
		Options:  outputFormats,
		Default:  "detailed",
		PageSize: tuiPageSize,
	}
	return survey.AskOne(formatPrompt, &selections.OutputFormat, survey.WithValidator(survey.Required))
}

// selectRegressionThreshold handles regression threshold selection
func selectRegressionThreshold(selections *tracker.Selections) error {
	thresholdPrompt := &survey.Confirm{
		Message: "Do you want to fail on performance regressions?",
		Default: false,
	}
	if err := survey.AskOne(thresholdPrompt, &selections.UseThreshold); err != nil {
		return err
	}

	if selections.UseThreshold {
		var thresholdStr string
		thresholdInputPrompt := &survey.Input{
			Message: "Enter regression threshold percentage (e.g., 5.0 for 5%):",
			Default: "5.0",
		}
		if err := survey.AskOne(thresholdInputPrompt, &thresholdStr, survey.WithValidator(survey.Required)); err != nil {
			return err
		}
		threshold, convErr := strconv.ParseFloat(thresholdStr, 64)
		if convErr != nil || threshold <= 0 {
			return fmt.Errorf("invalid threshold: %s", thresholdStr)
		}
		selections.RegressionThreshold = threshold
	}

	return nil
}

// setGlobalTrackingVariables sets the global CLI variables for tracking
func setGlobalTrackingVariables(selections *tracker.Selections) {
	Baseline = selections.Baseline
	Current = selections.Current
	benchmarkName = selections.BenchmarkName
	profileType = selections.ProfileType
	outputFormat = selections.OutputFormat
	failOnRegression = selections.UseThreshold
	regressionThreshold = selections.RegressionThreshold
}

func createTuiCmd() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "tui",
		Short: "Interactive selection of benchmarks and profiles, then runs prof auto",
		RunE:  runTUI,
	}

	cmd.AddCommand(createTuiTrackAutoCmd())

	return cmd
}

func createTuiTrackAutoCmd() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "track",
		Short: "Interactive tracking with existing benchmark data",
		RunE:  runTUITrackAuto,
	}

	return cmd
}

// discoverAvailableTags scans the bench directory for existing tags
func discoverAvailableTags() ([]string, error) {
	root, err := internal.FindGoModuleRoot()
	if err != nil {
		return nil, fmt.Errorf("failed to locate module root: %w", err)
	}

	benchDir := filepath.Join(root, internal.MainDirOutput)
	entries, err := os.ReadDir(benchDir)
	if err != nil {
		if os.IsNotExist(err) {
			return []string{}, nil
		}
		return nil, fmt.Errorf("failed to read bench directory: %w", err)
	}

	var tags []string
	for _, entry := range entries {
		if entry.IsDir() {
			tags = append(tags, entry.Name())
		}
	}

	return tags, nil
}

// discoverAvailableBenchmarks scans a specific tag directory for available benchmarks
func discoverAvailableBenchmarks(tag string) ([]string, error) {
	root, err := internal.FindGoModuleRoot()
	if err != nil {
		return nil, fmt.Errorf("failed to locate module root: %w", err)
	}

	benchDir := filepath.Join(root, internal.MainDirOutput, tag, internal.ProfileTextDir)
	entries, err := os.ReadDir(benchDir)
	if err != nil {
		if os.IsNotExist(err) {
			return []string{}, nil
		}
		return nil, fmt.Errorf("failed to read benchmark directory for tag %s: %w", tag, err)
	}

	var availableBenchmarks []string
	for _, entry := range entries {
		if entry.IsDir() {
			availableBenchmarks = append(availableBenchmarks, entry.Name())
		}
	}

	return availableBenchmarks, nil
}

// discoverAvailableProfiles scans a specific tag and benchmark for available profile types
func discoverAvailableProfiles(tag, benchmarkName string) ([]string, error) {
	root, err := internal.FindGoModuleRoot()
	if err != nil {
		return nil, fmt.Errorf("failed to locate module root: %w", err)
	}

	benchDir := filepath.Join(root, internal.MainDirOutput, tag, internal.ProfileTextDir, benchmarkName)
	entries, err := os.ReadDir(benchDir)
	if err != nil {
		if os.IsNotExist(err) {
			return []string{}, nil
		}
		return nil, fmt.Errorf("failed to read profile directory for tag %s, benchmark %s: %w", tag, benchmarkName, err)
	}

	var availableProfiles []string
	for _, entry := range entries {
		if !entry.IsDir() && strings.HasSuffix(entry.Name(), ".txt") {
			// Extract profile type from filename like "BenchmarkName_cpu.txt"
			name := entry.Name()
			if strings.HasPrefix(name, benchmarkName+"_") {
				profileTypeName := strings.TrimSuffix(strings.TrimPrefix(name, benchmarkName+"_"), ".txt")
				if profileTypeName == "cpu" || profileTypeName == "memory" || profileTypeName == "mutex" || profileTypeName == "block" {
					availableProfiles = append(availableProfiles, profileTypeName)
				}
			}
		}
	}

	return availableProfiles, nil
}

func runTUI(_ *cobra.Command, _ []string) error {
	// Get current working directory for scope-aware benchmark discovery
	currentDir, err := os.Getwd()
	if err != nil {
		return fmt.Errorf("failed to get current working directory: %w", err)
	}

	benchNames, err := benchmark.DiscoverBenchmarks(currentDir)
	if err != nil {
		return fmt.Errorf("failed to discover benchmarks: %w", err)
	}

	if len(benchNames) == 0 {
		return errors.New("no benchmarks found in this directory or its subdirectories (look for func BenchmarkXxx(b *testing.B) in *_test.go)")
	}

	var selectedBenches []string
	benchPrompt := &survey.MultiSelect{
		Message:  "Select benchmarks to run:",
		Options:  benchNames,
		PageSize: tuiPageSize,
	}
	if err = survey.AskOne(benchPrompt, &selectedBenches, survey.WithValidator(survey.Required)); err != nil {
		return err
	}

	profilesOptions := benchmark.SupportedProfiles
	var selectedProfiles []string
	profilesPrompt := &survey.MultiSelect{
		Message: "Select profiles:",
		Options: profilesOptions,
		Default: []string{"cpu"},
	}

	if err = survey.AskOne(profilesPrompt, &selectedProfiles, survey.WithValidator(survey.Required)); err != nil {
		return err
	}

	var countStr string
	countPrompt := &survey.Input{Message: "Number of runs (count):", Default: "1"}
	if err = survey.AskOne(countPrompt, &countStr, survey.WithValidator(survey.Required)); err != nil {
		return err
	}
	runCount, convErr := strconv.Atoi(countStr)
	if convErr != nil || runCount < 1 {
		return fmt.Errorf("invalid count: %s", countStr)
	}

	var tagStr string
	tagPrompt := &survey.Input{Message: "Tag name (used to group results under bench/<tag>):"}
	if err = survey.AskOne(tagPrompt, &tagStr, survey.WithValidator(survey.Required)); err != nil {
		return err
	}

	if err = benchmark.RunBenchmarks(selectedBenches, selectedProfiles, tagStr, runCount); err != nil {
		return err
	}

	return nil
}

func runTUITrackAuto(_ *cobra.Command, _ []string) error {
	// Discover available tags
	tags, err := discoverAvailableTags()
	if err != nil {
		return fmt.Errorf("failed to discover available tags: %w", err)
	}
	if len(tags) < minTagsForComparison {
		return errors.New("need at least 2 tags to compare (run 'prof tui' first to collect some data)")
	}

	// Get user selections
	selections, err := getTrackSelections(tags)
	if err != nil {
		return err
	}

	// Set global variables for the existing tracking logic
	setGlobalTrackingVariables(selections)

	// Now run the actual tracking command
	fmt.Printf("\n🚀 Running: prof track auto --base %s --current %s --bench-name %s --profile-type %s --output-format %s",
		selections.Baseline, selections.Current, selections.BenchmarkName, selections.ProfileType, selections.OutputFormat)
	if selections.UseThreshold {
		fmt.Printf(" --fail-on-regression --regression-threshold %.1f", selections.RegressionThreshold)
	}
	fmt.Println()

	return tracker.RunTrackAuto(selections)
}