package main

import (
	"fmt"
	"io"
	"log/slog"
	"math"
	"net/http"
	"os"
	"regexp"
	"strconv"
	"strings"
)

type cpuHandler struct {
	corePaths []string
	maxFreq   int
	minFreq   int
	logger    *slog.Logger
}

type memHandler struct {
	amount int64
	logger *slog.Logger
}

// readIntFromFile reads the given file until the end
// and converts the contained string into an integer
func readIntFromFile(path string) (int, error) {
	b, err := os.ReadFile(path)
	if err != nil {
		return 0, err
	}
	return strconv.Atoi(strings.TrimSpace(string(b)))
}

func readTotalMemoryKb() (int, error) {
	b, err := os.ReadFile("/proc/meminfo")
	if err != nil {
		return 0, err
	}
	s := string(b)
	for line := range strings.Lines(s) {
		if !strings.Contains(line, "MemTotal:") {
			continue
		}
		fields := strings.Fields(line)
		if len(fields) < 2 {
			return 0, fmt.Errorf("expected line '%s' to have at least 2 fields", line)
		}
		return strconv.Atoi(fields[1])
	}
	return 0, fmt.Errorf("failed to find MemTotal in /proc/meminfo")
}

func enumerateCPUCores(logger *slog.Logger) ([]string, int, int, error) {
	files, err := os.ReadDir("/sys/devices/system/cpu/")
	if err != nil {
		logger.Error("reading /sys/devices/system/cpu/", "err", err)
		return nil, -1, -1, err
	}
	r, _ := regexp.Compile("cpu([0-9]+)")

	freqPaths := make([]string, 0)
	maxF := math.MaxInt
	minF := math.MinInt

	// we look for cpuN directories
	for _, f := range files {
		if !f.IsDir() || !r.MatchString(f.Name()) {
			continue
		}

		// read files for min & max freqs
		dir := "/sys/devices/system/cpu/" + f.Name() + "/cpufreq/"
		minFreq, err := readIntFromFile(dir + "cpuinfo_min_freq")
		if err != nil {
			logger.Error("failed to read cpu min freq", "path", dir+"cpuinfo_min_freq", "error", err)
			continue
		}
		maxFreq, err := readIntFromFile(dir + "cpuinfo_max_freq")
		if err != nil {
			logger.Error("failed to read cpu max freq", "path", dir+"cpuinfo_min_freq", "error", err)
			continue
		}
		freqPaths = append(freqPaths, dir+"scaling_max_freq")
		maxF = min(maxF, maxFreq)
		minF = max(minF, minFreq)
	}
	return freqPaths, minF, maxF, nil
}

func (s *cpuHandler) handleMin(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" {
		writeResponse("invalid method", http.StatusBadRequest, w, s.logger)
		return
	}
	writeResponse(strconv.Itoa(s.minFreq), http.StatusOK, w, s.logger)
}

func (s *cpuHandler) handleMax(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" {
		writeResponse("invalid method", http.StatusBadRequest, w, s.logger)
		return
	}
	writeResponse(strconv.Itoa(s.maxFreq), http.StatusOK, w, s.logger)
}

func (s *cpuHandler) handleCurrent(w http.ResponseWriter, r *http.Request) {
	switch r.Method {
	case "GET":
		maxFreq := math.MinInt
		for _, core := range s.corePaths {
			freq, err := readIntFromFile(core)
			if err != nil {
				s.logger.Error("error while reading core frequency", "core", core, "error", err)
				writeResponse("internal server error", http.StatusInternalServerError, w, s.logger)
				return
			}
			s.logger.Debug("core frequency", "core", core, "frequency", freq)
			maxFreq = max(maxFreq, freq)
		}
		s.logger.Debug("reported frequency", "freq", maxFreq)
		writeResponse(strconv.Itoa(maxFreq), http.StatusOK, w, s.logger)
	case "PUT":
		body, err := io.ReadAll(r.Body)
		if err != nil {
			s.logger.Error("failed to read the request body", "err", err)
			writeResponse("internal server error", http.StatusInternalServerError, w, s.logger)
		}
		bodyStr := strings.TrimSpace(string(body))
		// Check if it is an integer
		freqInt, err := strconv.Atoi(bodyStr)
		if err != nil {
			s.logger.Error("failed to convert body to int", "error", err)
			writeResponse("bad request", http.StatusBadRequest, w, s.logger)
		}
		if freqInt < s.minFreq || freqInt > s.maxFreq {
			s.logger.Warn("requested frequency is out-of-bounds", "freq", freqInt, "min", s.minFreq, "max", s.maxFreq)
			freqInt = max(s.minFreq, min(freqInt, s.maxFreq))
			bodyStr = strconv.Itoa(freqInt)
		}
		s.logger.Info("requested to set current frequency", "frequency", bodyStr)
		success := true
		for _, core := range s.corePaths {
			err = os.WriteFile(core, []byte(bodyStr), 0)
			if err != nil {
				s.logger.Error("failed to set core frequency", "core", core, "error", err)
				success = false
			}
		}
		if success {
			w.WriteHeader(http.StatusOK)
		} else {
			w.WriteHeader(http.StatusInternalServerError)
		}
	default:
		s.logger.Error("unexpected method for current cpu", "method", r.Method)
		writeResponse("invalid method", http.StatusBadRequest, w, s.logger)
	}
}

func (s *cpuHandler) handleCount(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" {
		s.logger.Error("unexepcted method for cpu count", "method", r.Method)
		writeResponse("invalid method", http.StatusBadRequest, w, s.logger)
		return
	}

	count := len(s.corePaths)
	writeResponse(strconv.Itoa(count), http.StatusOK, w, s.logger)
}

func (s *cpuHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	requestedPath := getStrippedRequestPath(r.URL, r.Pattern)
	switch requestedPath {
	case "frequency/min":
		s.handleMin(w, r)
		return
	case "frequency/max":
		s.handleMax(w, r)
		return
	case "frequency/current":
		s.handleCurrent(w, r)
		return
	case "count":
		s.handleCount(w, r)
		return
	default:
		s.logger.Error("unknown resource", "path", requestedPath)
		writeResponse("not found", http.StatusNotFound, w, s.logger)
	}
}

func (m *memHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	requestedPath := getStrippedRequestPath(r.URL, r.Pattern)
	if requestedPath == "amount" {
		writeResponse(strconv.FormatInt(m.amount, 10), http.StatusOK, w, m.logger)
	} else {
		m.logger.Error("unknown resource", "path", requestedPath)
		writeResponse("not found", http.StatusNotFound, w, m.logger)
	}
}

func main() {
	appEnv := os.Getenv("AGENT_ENV")
	loggerOpts := &slog.HandlerOptions{
		Level: slog.LevelDebug,
	}
	var handler slog.Handler = slog.NewTextHandler(os.Stdout, loggerOpts)
	if appEnv == "production" {
		loggerOpts.Level = slog.LevelInfo
		handler = slog.NewJSONHandler(os.Stdout, loggerOpts)
	}
	logger := slog.New(handler)

	freqPaths, minF, maxF, err := enumerateCPUCores(logger)
	if err != nil {
		os.Exit(1)
	}

	cpuHandler := &cpuHandler{
		logger:    logger,
		maxFreq:   maxF,
		minFreq:   minF,
		corePaths: freqPaths,
	}

	memKb, err := readTotalMemoryKb()
	if err != nil {
		logger.Error("failed to read memory size", "err", err.Error())
		os.Exit(1)
	}
	memHandler := &memHandler{
		amount: int64(memKb) * 1024,
		logger: logger,
	}

	http.Handle("/cpu/", logRequests(logger, cpuHandler))
	http.Handle("/mem/", logRequests(logger, memHandler))
	logger.Info("exit", "result", http.ListenAndServe(":8080", nil))
}