package drivers

import (
	"fmt"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"
)

const milliJoulePerJoule = float64(1000000)

type EnergyState struct {
	lastEnergy uint64
	lastTime   time.Time
}

type RaplDriver struct {
	energyState map[string]EnergyState
}

func NewRaplDriver() *RaplDriver {
	return &RaplDriver{
		energyState: map[string]EnergyState{},
	}

}

func (r *RaplDriver) ReadPowerSensors() ([]PowerSensor, error) {
	powerSensors := []PowerSensor{}

	patterns := []string{
		"/sys/class/powercap/intel-rapl*/energy_uj",
		"/sys/class/powercap/intel-rapl/intel-rapl*/energy_uj",
		"/sys/class/powercap/intel-rapl/intel-rapl*/*/energy_uj",
		"/sys/class/powercap/intel-rapl:*/energy_uj",
		"/sys/class/powercap/intel-rapl:*/*/energy_uj",
	}

	energyInputs := []string{}
	unique := make(map[string]bool)

	for _, pattern := range patterns {
		matches, _ := filepath.Glob(pattern)

		for _, match := range matches {
			realPath, err := filepath.EvalSymlinks(match)
			if err != nil {
				realPath = match
			}
			if !unique[realPath] {
				unique[realPath] = true
				energyInputs = append(energyInputs, realPath)
			}
		}
	}

	if len(energyInputs) == 0 {
		return powerSensors, fmt.Errorf("found no compatible rapl drivers")
	}

	r.calculateEnergyDelta(energyInputs, &powerSensors)
	return powerSensors, nil
}

func getRaplLabel(inputPath string) string {
	dir := filepath.Dir(inputPath)
	if nameBytes, err := os.ReadFile(filepath.Join(dir, "name")); err == nil {
		return strings.TrimSpace(string(nameBytes))
	}
	return filepath.Base(dir)
}

func (r *RaplDriver) calculateEnergyDelta(inputs []string, out *[]PowerSensor) {
	for _, path := range inputs {
		now := time.Now()
		valBytes, err := os.ReadFile(path)
		if err != nil {
			continue
		}

		microJoules, err := strconv.ParseUint(strings.TrimSpace(string(valBytes)), 10, 64)
		if err != nil {
			continue
		}

		sensorLabel := getRaplLabel(path)

		prevState, exists := r.energyState[path]

		if !exists {

			r.energyState[path] = EnergyState{
				lastEnergy: microJoules,
				lastTime:   now,
			}
			*out = append(*out, PowerSensor{
				Name:  sensorLabel,
				Watts: 0.0,
			})
			continue
		}

		timeDiff := now.Sub(prevState.lastTime).Seconds()

		// technically the energy counter could overflow, at ~18 TerraJoules
		// the value will be off for that one measurement, then become correct again
		energyDiff := float64(microJoules - prevState.lastEnergy)

		watts := (energyDiff / milliJoulePerJoule) / timeDiff

		r.energyState[path] = EnergyState{
			lastEnergy: microJoules,
			lastTime:   now,
		}

		*out = append(*out, PowerSensor{
			Name:  sensorLabel,
			Watts: watts,
		})
	}
}