package main

import (
	"bufio"
	"errors"
	"fmt"
	"os/exec"
	"strconv"
	"strings"
)

// Custom error types
var (
	ErrParse   = errors.New("socinfo parsing error")
	ErrCommand = errors.New("command execution error")
	ErrConvert = errors.New("conversion error")
)

type (
	Watts     = uint32
	Bandwidth = uint32
	CoreCount = uint16
)

// SocInfo represents information about the Silicon chip
type SocInfo struct {
	CPUBrandName string     `json:"cpu_brand_name"`
	NumCPUCores  CoreCount  `json:"num_cpu_cores"`
	NumGPUCores  CoreCount  `json:"num_gpu_cores"`
	CPUMaxPower  *Watts     `json:"cpu_max_power,omitempty"`
	GPUMaxPower  *Watts     `json:"gpu_max_power,omitempty"`
	CPUMaxBW     *Bandwidth `json:"cpu_max_bw,omitempty"`
	GPUMaxBW     *Bandwidth `json:"gpu_max_bw,omitempty"`
	ECoreCount   CoreCount  `json:"e_core_count"`
	PCoreCount   CoreCount  `json:"p_core_count"`
}

// AppleChip represents different Apple Silicon variants
type AppleChip int

const (
	M1 AppleChip = iota
	M1Pro
	M1Max
	M1Ultra
	M2
	M2Pro
	M2Max
	M2Ultra
	M3
	M3Pro
	M3Max
	Unknown
)

// ChipSpecs holds the specifications for each chip variant
type ChipSpecs struct {
	CPUTDP Watts
	GPUTDP Watts
	CPUBW  Bandwidth
	GPUBW  Bandwidth
}

// fromBrandString determines the Apple chip type from brand string
func fromBrandString(brand string) AppleChip {
	switch {
	case strings.Contains(brand, "M1 Pro"):
		return M1Pro
	case strings.Contains(brand, "M1 Max"):
		return M1Max
	case strings.Contains(brand, "M1 Ultra"):
		return M1Ultra
	case strings.Contains(brand, "M1"):
		return M1
	case strings.Contains(brand, "M2 Pro"):
		return M2Pro
	case strings.Contains(brand, "M2 Max"):
		return M2Max
	case strings.Contains(brand, "M2 Ultra"):
		return M2Ultra
	case strings.Contains(brand, "M2"):
		return M2
	case strings.Contains(brand, "M3 Pro"):
		return M3Pro
	case strings.Contains(brand, "M3 Max"):
		return M3Max
	case strings.Contains(brand, "M3"):
		return M3
	default:
		return Unknown
	}
}

// getSpecs returns the specifications for the chip
func (chip AppleChip) getSpecs() ChipSpecs {
	switch chip {
	case M1:
		return ChipSpecs{CPUTDP: 20, GPUTDP: 20, CPUBW: 70, GPUBW: 70}
	case M1Pro:
		return ChipSpecs{CPUTDP: 30, GPUTDP: 30, CPUBW: 200, GPUBW: 200}
	case M1Max:
		return ChipSpecs{CPUTDP: 30, GPUTDP: 60, CPUBW: 250, GPUBW: 400}
	case M1Ultra:
		return ChipSpecs{CPUTDP: 60, GPUTDP: 120, CPUBW: 500, GPUBW: 800}
	case M2:
		return ChipSpecs{CPUTDP: 25, GPUTDP: 15, CPUBW: 100, GPUBW: 100}
	case M2Pro:
		return ChipSpecs{CPUTDP: 30, GPUTDP: 35, CPUBW: 0, GPUBW: 0}
	case M2Max:
		return ChipSpecs{CPUTDP: 30, GPUTDP: 40, CPUBW: 0, GPUBW: 0}
	default:
		return ChipSpecs{CPUTDP: 0, GPUTDP: 0, CPUBW: 0, GPUBW: 0}
	}
}

// SystemCommand interface for testable command execution
type SystemCommand interface {
	Execute(binary string, args ...string) ([]byte, error)
}

// RealCommand implements SystemCommand for actual system calls
type RealCommand struct{}

func (r *RealCommand) Execute(binary string, args ...string) ([]byte, error) {
	cmd := exec.Command(binary, args...)
	output, err := cmd.Output()
	if err != nil {
		return nil, fmt.Errorf("%w: %v", ErrCommand, err)
	}
	return output, nil
}

const sysctlPath = "/usr/sbin/sysctl"

// getCPUInfo retrieves CPU information using sysctl
func getCPUInfo(cmd SystemCommand) (string, CoreCount, CoreCount, CoreCount, error) {
	args := []string{
		"-n", // don't display the variable name
		"machdep.cpu.brand_string",
		"machdep.cpu.core_count",
		"hw.perflevel0.logicalcpu",
		"hw.perflevel1.logicalcpu",
	}

	output, err := cmd.Execute(sysctlPath, args...)
	if err != nil {
		return "", 0, 0, 0, err
	}

	lines := strings.Split(strings.TrimSpace(string(output)), "\n")
	if len(lines) < 4 {
		return "", 0, 0, 0, fmt.Errorf("%w: insufficient output lines", ErrParse)
	}

	cpuBrandName := lines[0]

	numCPUCores, err := strconv.ParseUint(lines[1], 10, 16)
	if err != nil {
		return "", 0, 0, 0, fmt.Errorf("%w: parsing cpu cores: %v", ErrConvert, err)
	}

	numPerformanceCores, err := strconv.ParseUint(lines[2], 10, 16)
	if err != nil {
		return "", 0, 0, 0, fmt.Errorf("%w: parsing performance cores: %v", ErrConvert, err)
	}

	numEfficiencyCores, err := strconv.ParseUint(lines[3], 10, 16)
	if err != nil {
		return "", 0, 0, 0, fmt.Errorf("%w: parsing efficiency cores: %v", ErrConvert, err)
	}

	return cpuBrandName, CoreCount(
			numCPUCores,
		), CoreCount(
			numPerformanceCores,
		), CoreCount(
			numEfficiencyCores,
		), nil
}

// getGPUInfo retrieves GPU information using system_profiler
func getGPUInfo(cmd SystemCommand) (CoreCount, error) {
	args := []string{"-detailLevel", "basic", "SPDisplaysDataType"}

	output, err := cmd.Execute("/usr/sbin/system_profiler", args...)
	if err != nil {
		return 0, err
	}

	scanner := bufio.NewScanner(strings.NewReader(string(output)))
	for scanner.Scan() {
		line := scanner.Text()
		trimmed := strings.TrimSpace(line)
		if strings.HasPrefix(trimmed, "Total Number of Cores") {
			parts := strings.Split(trimmed, ": ")
			if len(parts) == 2 {
				cores, err := strconv.ParseUint(strings.TrimSpace(parts[1]), 10, 16)
				if err != nil {
					return 0, fmt.Errorf("%w: parsing gpu cores: %v", ErrConvert, err)
				}
				return CoreCount(cores), nil
			}
		}
	}

	if err := scanner.Err(); err != nil {
		return 0, fmt.Errorf("%w: scanning output: %v", ErrParse, err)
	}

	return 0, fmt.Errorf("%w: GPU core count not found", ErrParse)
}

// NewSocInfo creates a new SocInfo instance by gathering system information
func NewSocInfo() (*SocInfo, error) {
	cmd := &RealCommand{}
	return NewSocInfoWithCommand(cmd)
}

// NewSocInfoWithCommand creates a new SocInfo instance with a custom command executor (for testing)
func NewSocInfoWithCommand(cmd SystemCommand) (*SocInfo, error) {
	cpuBrandName, numCPUCores, eCoreCount, pCoreCount, err := getCPUInfo(cmd)
	if err != nil {
		return nil, err
	}

	numGPUCores, err := getGPUInfo(cmd)
	if err != nil {
		return nil, err
	}

	chip := fromBrandString(cpuBrandName)
	specs := chip.getSpecs()

	var cpuMaxPower, gpuMaxPower *Watts
	var cpuMaxBW, gpuMaxBW *Bandwidth

	if specs.CPUTDP > 0 {
		cpuMaxPower = &specs.CPUTDP
	}
	if specs.GPUTDP > 0 {
		gpuMaxPower = &specs.GPUTDP
	}
	if specs.CPUBW > 0 {
		cpuMaxBW = &specs.CPUBW
	}
	if specs.GPUBW > 0 {
		gpuMaxBW = &specs.GPUBW
	}

	return &SocInfo{
		CPUBrandName: cpuBrandName,
		NumCPUCores:  numCPUCores,
		NumGPUCores:  numGPUCores,
		CPUMaxPower:  cpuMaxPower,
		GPUMaxPower:  gpuMaxPower,
		CPUMaxBW:     cpuMaxBW,
		GPUMaxBW:     gpuMaxBW,
		ECoreCount:   eCoreCount,
		PCoreCount:   pCoreCount,
	}, nil
}

func main() {
	socInfo, err := NewSocInfo()
	if err != nil {
		fmt.Printf("Error getting SoC info: %v\n", err)
		return
	}

	var cpuPower Watts
	if socInfo.CPUMaxPower != nil {
		cpuPower = *socInfo.CPUMaxPower
	}

	fmt.Printf("Our CPU is an %s, and we have %d CPU cores, and %d GPU cores. The TDP is %d.\n",
		socInfo.CPUBrandName,
		socInfo.NumCPUCores,
		socInfo.NumGPUCores,
		cpuPower,
	)
}