package node

import (
	"bytes"
	"context"
	"fmt"
	"k8s.io/klog/v2"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"
	"time"
)

func ConfigureUBootCommands(ctx context.Context, n *NodeContext) error {
	const (
		procMTDPath  = "/proc/mtd"
		fwEnvCfgPath = "/etc/fw_env.config"
		targetName   = "uboot-env"
	)

	_ = n

	type mtdInfo struct {
		dev       string
		sizeHex   string
		eraseHex  string
		name      string
		partSize  uint64
		eraseSize uint64
	}

	parseHex := func(s string) (uint64, error) {
		var v uint64
		_, err := fmt.Sscanf(s, "%x", &v)
		if err != nil {
			return 0, fmt.Errorf("parse hex %q: %w", s, err)
		}
		return v, nil
	}

	raw, err := os.ReadFile(procMTDPath)
	if err != nil {
		return fmt.Errorf("read %s: %w", procMTDPath, err)
	}

	lineRE := regexp.MustCompile(`^(mtd\d+):\s+([0-9a-fA-F]+)\s+([0-9a-fA-F]+)\s+"([^"]+)"$`)

	var envMTD *mtdInfo
	for _, line := range strings.Split(string(raw), "\n") {
		klog.V(4).Info(line)
		line = strings.TrimSpace(line)
		if line == "" || strings.HasPrefix(line, "dev:") {
			continue
		}
		m := lineRE.FindStringSubmatch(line)
		if m == nil {
			continue
		}
		if m[4] != targetName {
			continue
		}

		partSize, err := parseHex(m[2])
		if err != nil {
			return fmt.Errorf("parse partition size for %s: %w", m[1], err)
		}
		eraseSize, err := parseHex(m[3])
		if err != nil {
			return fmt.Errorf("parse erase size for %s: %w", m[1], err)
		}

		envMTD = &mtdInfo{
			dev:       "/dev/" + m[1],
			sizeHex:   "0x" + strings.ToLower(m[2]),
			eraseHex:  "0x" + strings.ToLower(m[3]),
			name:      m[4],
			partSize:  partSize,
			eraseSize: eraseSize,
		}
		break
	}

	if envMTD == nil {
		return fmt.Errorf("no %q partition found in %s", targetName, procMTDPath)
	}

	if _, err := exec.LookPath("fw_printenv"); err != nil {
		return fmt.Errorf("fw_printenv not found in PATH: %w", err)
	}

	// Candidate env sizes to probe, smallest first.
	// Many boards use 0x2000/0x4000/0x8000/0x10000.
	// Keep candidates <= partition size and > 0.
	candidateSizes := uniqueUint64s([]uint64{
		0x2000,
		0x4000,
		0x8000,
		0x10000,
		envMTD.eraseSize,
	})

	var filtered []uint64
	for _, sz := range candidateSizes {
		if sz > 0 && sz <= envMTD.partSize {
			filtered = append(filtered, sz)
		}
	}
	if len(filtered) == 0 {
		return fmt.Errorf("no valid candidate env sizes for %s", envMTD.dev)
	}

	type probeResult struct {
		size   uint64
		config string
		output string
	}

	var best *probeResult
	var attempts []probeResult

	for _, sz := range filtered {
		cfg := fmt.Sprintf("%s 0x0 0x%x 0x%x\n", envMTD.dev, sz, envMTD.eraseSize)
		klog.Infof("Trying: %s", cfg)

		tmpDir, err := os.MkdirTemp("", "fwenv-probe-*")
		if err != nil {
			return fmt.Errorf("create temp dir: %w", err)
		}
		cfgPath := filepath.Join(tmpDir, "fw_env.config")
		if err := os.WriteFile(cfgPath, []byte(cfg), 0o644); err != nil {
			_ = os.RemoveAll(tmpDir)
			return fmt.Errorf("write temp config: %w", err)
		}

		out, _ := runFWPrintenvWithConfig(ctx, cfgPath)
		_ = os.RemoveAll(tmpDir)

		res := probeResult{
			size:   sz,
			config: cfg,
			output: out,
		}
		attempts = append(attempts, res)

		if looksLikeValidUBootEnv(out) {
			best = &res
			break
		}
	}

	if best == nil {
		var b strings.Builder
		b.WriteString("could not determine correct fw_env.config by probing candidate sizes\n")
		b.WriteString(fmt.Sprintf("partition: %s size=%s erase=%s\n", envMTD.dev, envMTD.sizeHex, envMTD.eraseHex))
		for _, a := range attempts {
			b.WriteString(fmt.Sprintf("\n--- candidate env-size=0x%x ---\n", a.size))
			b.WriteString(trimOutput(a.output, 600))
			b.WriteString("\n")
		}
		return fmt.Errorf(b.String())
	}

	klog.Infof("Using: %s", best.config)

	// Avoid rewriting if already identical.
	existing, err := os.ReadFile(fwEnvCfgPath)
	if err == nil && bytes.Equal(existing, []byte(best.config)) {
		return nil
	}
	if err != nil && !os.IsNotExist(err) {
		return fmt.Errorf("read %s: %w", fwEnvCfgPath, err)
	}

	if err := os.MkdirAll(filepath.Dir(fwEnvCfgPath), 0o755); err != nil {
		return fmt.Errorf("mkdir %s: %w", filepath.Dir(fwEnvCfgPath), err)
	}
	if err := os.WriteFile(fwEnvCfgPath, []byte(best.config), 0o644); err != nil {
		return fmt.Errorf("write %s: %w", fwEnvCfgPath, err)
	}

	return nil
}

func runFWPrintenvWithConfig(ctx context.Context, configPath string) (string, error) {
	ctx, cancel := context.WithTimeout(ctx, 3*time.Second)
	defer cancel()

	cmd := exec.CommandContext(ctx, "fw_printenv", "-c", configPath)
	out, err := cmd.CombinedOutput()
	return string(out), err
}

func looksLikeValidUBootEnv(out string) bool {
	lower := strings.ToLower(out)

	// Hard fail patterns.
	if strings.Contains(lower, "bad crc") {
		return false
	}
	if strings.Contains(lower, "using default environment") {
		return false
	}
	if strings.Contains(lower, "cannot parse config file") {
		return false
	}
	if strings.Contains(lower, "failed to find nvmem device") {
		return false
	}
	if strings.Contains(lower, "configuration file wrong") {
		return false
	}

	// Need a few plausible env variables.
	lines := strings.Split(out, "\n")
	var hits int
	for _, line := range lines {
		line = strings.TrimSpace(line)
		if !strings.Contains(line, "=") {
			continue
		}
		switch {
		case strings.HasPrefix(line, "bootcmd="),
			strings.HasPrefix(line, "bootdelay="),
			strings.HasPrefix(line, "baudrate="),
			strings.HasPrefix(line, "bootargs="),
			strings.HasPrefix(line, "altbootcmd="),
			strings.HasPrefix(line, "stderr="),
			strings.HasPrefix(line, "stdin="),
			strings.HasPrefix(line, "stdout="):
			hits++
		}
	}

	return hits >= 2
}

func uniqueUint64s(in []uint64) []uint64 {
	seen := make(map[uint64]struct{}, len(in))
	out := make([]uint64, 0, len(in))
	for _, v := range in {
		if _, ok := seen[v]; ok {
			continue
		}
		seen[v] = struct{}{}
		out = append(out, v)
	}
	return out
}

func trimOutput(s string, max int) string {
	s = strings.TrimSpace(s)
	if len(s) <= max {
		return s
	}
	return s[:max] + "...(truncated)"
}