package bootenv

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

const (
	// RPi partition numbers: slot A = partition 2, slot B = partition 3.
	rpiSlotAPartition = 2
	rpiSlotBPartition = 3

	defaultBootCounter = 3
)

// RPiEnv implements BootEnv using Raspberry Pi firmware autoboot.txt.
type RPiEnv struct {
	autobootPath string // path to autoboot.txt
	statusPath   string // path to boot-status file
}

// NewRPi creates a new RPi-based BootEnv.
// dir is the directory containing autoboot.txt (typically the boot control
// partition mount point).
func NewRPi(dir string) BootEnv {
	return &RPiEnv{
		autobootPath: filepath.Join(dir, "autoboot.txt"),
		statusPath:   filepath.Join(dir, "boot-status"),
	}
}

func (r *RPiEnv) ActiveSlot() (string, error) {
	partNum, err := r.readAllBootPartition()
	if err != nil {
		return "", fmt.Errorf("reading active slot: %w", err)
	}
	return partNumToSlot(partNum)
}

func (r *RPiEnv) PassiveSlot() (string, error) {
	active, err := r.ActiveSlot()
	if err != nil {
		return "", err
	}
	if active == SlotA {
		return SlotB, nil
	}
	return SlotA, nil
}

func (r *RPiEnv) BootCounter() (int, error) {
	status, err := r.readStatus()
	if err != nil {
		return -1, err
	}
	val, ok := status["boot_counter"]
	if !ok {
		return defaultBootCounter, nil
	}
	n, err := strconv.Atoi(val)
	if err != nil {
		return -1, fmt.Errorf("invalid boot_counter %q: %w", val, err)
	}
	return n, nil
}

func (r *RPiEnv) BootSuccess() (bool, error) {
	status, err := r.readStatus()
	if err != nil {
		return false, err
	}
	return status["boot_success"] == "1", nil
}

func (r *RPiEnv) MarkBootSuccess() error {
	// Make the current slot permanent by updating [all] boot_partition
	active, err := r.ActiveSlot()
	if err != nil {
		return fmt.Errorf("marking boot success: %w", err)
	}
	partNum := slotToPartNum(active)
	if err := r.writeAllBootPartition(partNum); err != nil {
		return err
	}
	return r.writeStatus(defaultBootCounter, true)
}

func (r *RPiEnv) ActivateSlot(slot string) error {
	if slot != SlotA && slot != SlotB {
		return fmt.Errorf("invalid slot: %q (must be A or B)", slot)
	}
	partNum := slotToPartNum(slot)
	// Update [tryboot] to point to the new slot
	if err := r.writeTrybootPartition(partNum); err != nil {
		return err
	}
	return r.writeStatus(defaultBootCounter, false)
}

func (r *RPiEnv) ForceRollback() error {
	passive, err := r.PassiveSlot()
	if err != nil {
		return err
	}
	// Swap the [all] boot_partition to the other slot
	partNum := slotToPartNum(passive)
	if err := r.writeAllBootPartition(partNum); err != nil {
		return err
	}
	if err := r.writeTrybootPartition(partNum); err != nil {
		return err
	}
	return r.writeStatus(defaultBootCounter, false)
}

// readAllBootPartition reads the boot_partition value from the [all] section.
func (r *RPiEnv) readAllBootPartition() (int, error) {
	sections, err := r.parseAutoboot()
	if err != nil {
		return 0, err
	}
	val, ok := sections["all"]["boot_partition"]
	if !ok {
		return 0, fmt.Errorf("boot_partition not found in [all] section")
	}
	return strconv.Atoi(val)
}

// writeAllBootPartition updates the [all] boot_partition value.
func (r *RPiEnv) writeAllBootPartition(partNum int) error {
	sections, err := r.parseAutoboot()
	if err != nil {
		return err
	}
	if sections["all"] == nil {
		sections["all"] = make(map[string]string)
	}
	sections["all"]["boot_partition"] = strconv.Itoa(partNum)
	return r.writeAutoboot(sections)
}

// writeTrybootPartition updates the [tryboot] boot_partition value.
func (r *RPiEnv) writeTrybootPartition(partNum int) error {
	sections, err := r.parseAutoboot()
	if err != nil {
		return err
	}
	if sections["tryboot"] == nil {
		sections["tryboot"] = make(map[string]string)
	}
	sections["tryboot"]["boot_partition"] = strconv.Itoa(partNum)
	return r.writeAutoboot(sections)
}

// parseAutoboot reads autoboot.txt into a map of section -> key=value pairs.
func (r *RPiEnv) parseAutoboot() (map[string]map[string]string, error) {
	data, err := os.ReadFile(r.autobootPath)
	if err != nil {
		return nil, fmt.Errorf("reading autoboot.txt: %w", err)
	}

	sections := make(map[string]map[string]string)
	currentSection := ""

	for _, line := range strings.Split(string(data), "\n") {
		line = strings.TrimSpace(line)
		if line == "" {
			continue
		}
		if strings.HasPrefix(line, "[") && strings.HasSuffix(line, "]") {
			currentSection = line[1 : len(line)-1]
			if sections[currentSection] == nil {
				sections[currentSection] = make(map[string]string)
			}
			continue
		}
		parts := strings.SplitN(line, "=", 2)
		if len(parts) == 2 && currentSection != "" {
			sections[currentSection][strings.TrimSpace(parts[0])] = strings.TrimSpace(parts[1])
		}
	}

	return sections, nil
}

// writeAutoboot writes sections back to autoboot.txt.
// Section order: [all] first, then [tryboot].
func (r *RPiEnv) writeAutoboot(sections map[string]map[string]string) error {
	var sb strings.Builder

	// Write [all] section first
	if all, ok := sections["all"]; ok {
		sb.WriteString("[all]\n")
		for k, v := range all {
			sb.WriteString(k + "=" + v + "\n")
		}
	}

	// Write [tryboot] section
	if tryboot, ok := sections["tryboot"]; ok {
		sb.WriteString("[tryboot]\n")
		for k, v := range tryboot {
			sb.WriteString(k + "=" + v + "\n")
		}
	}

	return os.WriteFile(r.autobootPath, []byte(sb.String()), 0o644)
}

// readStatus reads the boot-status key=value file.
func (r *RPiEnv) readStatus() (map[string]string, error) {
	data, err := os.ReadFile(r.statusPath)
	if err != nil {
		if os.IsNotExist(err) {
			// Return defaults if status file doesn't exist yet
			return map[string]string{
				"boot_counter": strconv.Itoa(defaultBootCounter),
				"boot_success": "0",
			}, nil
		}
		return nil, fmt.Errorf("reading boot-status: %w", err)
	}

	status := make(map[string]string)
	for _, line := range strings.Split(string(data), "\n") {
		line = strings.TrimSpace(line)
		if line == "" {
			continue
		}
		parts := strings.SplitN(line, "=", 2)
		if len(parts) == 2 {
			status[strings.TrimSpace(parts[0])] = strings.TrimSpace(parts[1])
		}
	}
	return status, nil
}

// writeStatus writes boot_counter and boot_success to the status file.
func (r *RPiEnv) writeStatus(counter int, success bool) error {
	successVal := "0"
	if success {
		successVal = "1"
	}
	content := fmt.Sprintf("boot_counter=%d\nboot_success=%s\n", counter, successVal)
	return os.WriteFile(r.statusPath, []byte(content), 0o644)
}

func partNumToSlot(partNum int) (string, error) {
	switch partNum {
	case rpiSlotAPartition:
		return SlotA, nil
	case rpiSlotBPartition:
		return SlotB, nil
	default:
		return "", fmt.Errorf("unknown partition number %d (expected %d or %d)", partNum, rpiSlotAPartition, rpiSlotBPartition)
	}
}

func slotToPartNum(slot string) int {
	if slot == SlotB {
		return rpiSlotBPartition
	}
	return rpiSlotAPartition
}