package main

import (
	"bufio"
	"context"
	"encoding/json"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"sort"
	"strconv"
	"strings"

	"golang.org/x/sync/errgroup"
)

type CPUInfo struct {
	CPU    int     `json:"cpu"`
	MaxMHz float64 `json:"maxmhz"`
}

type LSCPUOutput struct {
	CPUs []CPUInfo `json:"cpus"`
}

type CoreFreq struct {
	Core int
	Freq int
}

func main() {
	args := os.Args[1:]
	if len(args) == 0 {
		fmt.Println("Please provide a command to run")
		os.Exit(1)
	}

	eg, ctx := errgroup.WithContext(context.Background())
	canPerf, cores, err := getBestCores()
	if err != nil {
		eg.Go(func() error {
			return runCommand(ctx, false, []int{}, args)
		})
		err = eg.Wait()
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
		os.Exit(0)
	}

	eg.Go(func() error {
		return runCommand(ctx, canPerf, cores, args)
	})
	err = eg.Wait()
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
	os.Exit(0)
}

func runCommand(ctx context.Context, canPerf bool, cores []int, args []string) error {
	commandToRun := make([]string, 0)
	if canPerf {
		commandToRun = append(commandToRun, "powerprofilesctl", "launch", "-p", "performance")
	}

	if len(cores) > 0 {
		coreVals := make([]string, len(cores))
		for i, core := range cores {
			coreVals[i] = fmt.Sprint(core)
		}

		commandToRun = append(commandToRun, "taskset", "-c", strings.Join(coreVals, ","))
	}

	commandToRun = append(commandToRun, args...)

	cmd := exec.CommandContext(ctx, commandToRun[0], commandToRun[1:]...)
	cmd.Env = append(cmd.Environ(), "POSIXLY_CORRECT=1")
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	cmd.Stdin = os.Stdin

	return cmd.Run()
}

func getBestCores() (bool, []int, error) {
	perfMode, amdPrefCores, amdFlag, highestCores, err := GetPPStatusAndHighestCores()
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}

	if !amdFlag {
		return perfMode, highestCores, nil
	}

	numAmdCores := len(amdPrefCores)
	switch numAmdCores {
	case 32:
		return perfMode, getHighestCores(16, amdPrefCores), nil
	case 24:
		return perfMode, getHighestCores(12, amdPrefCores), nil
	case 16:
		return perfMode, getHighestCores(8, amdPrefCores), nil
	default:
		return perfMode, getHighestCores(numAmdCores, amdPrefCores), nil
	}
}

func getHighestCores(n int, coreFreqs []CoreFreq) []int {
	if n <= 0 {
		return []int{}
	}

	sort.Slice(coreFreqs, func(i, j int) bool {
		return coreFreqs[i].Freq > coreFreqs[j].Freq
	})

	result := make([]int, 0, n)
	for i := 0; i < n && i < len(coreFreqs); i++ {
		result = append(result, coreFreqs[i].Core)
	}
	sort.Stable(sort.IntSlice(result))

	return result
}

func GetPPStatusAndHighestCores() (bool, []CoreFreq, bool, []int, error) {
	var (
		performanceMode bool
		amdPstateFreqs  []CoreFreq
		amdFlag         bool
		highestCores    []int
	)

	g, ctx := errgroup.WithContext(context.Background())
	g.Go(func() error {
		cmd := exec.CommandContext(ctx, "powerprofilesctl", "list")
		output, err := cmd.Output()
		if err != nil {
			performanceMode = false
			return nil
		}

		performanceMode = strings.Contains(string(output), "performance")
		return nil
	})

	g.Go(func() error {
		pattern := "/sys/devices/system/cpu/cpu*/cpufreq/amd_pstate_prefcore_ranking"
		matches, err := filepath.Glob(pattern)
		if err != nil {
			return fmt.Errorf("failed to glob AMD pstate files: %w", err)
		}

		if len(matches) > 0 {
			amdFlag = true
		}

		for _, match := range matches {
			file, err := os.Open(match)
			if err != nil {
				return fmt.Errorf("failed to open file %s: %w", match, err)
			}
			defer file.Close()

			parts := strings.Split(match, "/")
			cpuPart := parts[len(parts)-3]
			coreNum, err := strconv.Atoi(strings.TrimPrefix(cpuPart, "cpu"))
			if err != nil {
				return fmt.Errorf("failed to parse core number from path %s: %w", match, err)
			}

			scanner := bufio.NewScanner(file)
			if scanner.Scan() {
				value, err := strconv.Atoi(strings.TrimSpace(scanner.Text()))
				if err != nil {
					return fmt.Errorf("failed to parse value from file %s: %w", match, err)
				}

				amdPstateFreqs = append(amdPstateFreqs, CoreFreq{Core: coreNum, Freq: value})
			}

			if err := scanner.Err(); err != nil {
				return fmt.Errorf("error reading file %s: %w", match, err)
			}
		}
		return nil
	})

	g.Go(func() error {
		cmd := exec.CommandContext(ctx, "lscpu", "-e", "-J")
		output, err := cmd.Output()
		if err != nil {
			return fmt.Errorf("failed to run lscpu command: %w", err)
		}

		var lscpuOutput LSCPUOutput
		if err := json.Unmarshal(output, &lscpuOutput); err != nil {
			return fmt.Errorf("failed to unmarshal lscpu output: %w", err)
		}

		var maxMHz float64
		for _, cpu := range lscpuOutput.CPUs {
			if cpu.MaxMHz > maxMHz {
				maxMHz = cpu.MaxMHz
			}
		}

		for _, cpu := range lscpuOutput.CPUs {
			if cpu.MaxMHz == maxMHz {
				highestCores = append(highestCores, cpu.CPU)
			}
		}

		return nil
	})

	if err := g.Wait(); err != nil {
		return false, nil, false, nil, err
	}

	return performanceMode, amdPstateFreqs, amdFlag, highestCores, nil
}