package cmd

import (
	"fmt"
	"log/slog"

	"github.com/portainer/kubesolo-os/update/pkg/image"
	"github.com/portainer/kubesolo-os/update/pkg/partition"
	"github.com/portainer/kubesolo-os/update/pkg/state"
)

// Apply downloads a new OS image and writes it to the passive partition.
// It does NOT activate the new partition — use 'activate' for that.
//
// State transitions: Idle/Success/Failed → Checking → Downloading → Staged.
// On any error the state moves to Failed with LastError set.
func Apply(args []string) error {
	opts := parseOpts(args)

	if opts.ServerURL == "" {
		return fmt.Errorf("--server is required")
	}

	st, err := state.Load(opts.StatePath)
	if err != nil {
		// Don't block the operation on a corrupt state file. Log + recover.
		slog.Warn("state file unreadable, starting fresh", "error", err)
		st = state.New()
	}

	env := opts.NewBootEnv()

	// Record the current running version as the "from" reference. The active
	// slot's version file is the most reliable source.
	activeSlot, slotErr := env.ActiveSlot()
	if slotErr == nil {
		if partInfo, perr := partition.GetSlotPartition(activeSlot); perr == nil {
			mp := "/tmp/kubesolo-active-" + activeSlot
			if merr := partition.MountReadOnly(partInfo.Device, mp); merr == nil {
				if v, rerr := partition.ReadVersion(mp); rerr == nil {
					st.SetFromVersion(v)
				}
				partition.Unmount(mp)
			}
		}
	}

	// Determine passive slot
	passiveSlot, err := env.PassiveSlot()
	if err != nil {
		_ = st.RecordError(opts.StatePath, fmt.Errorf("reading passive slot: %w", err))
		return fmt.Errorf("reading passive slot: %w", err)
	}

	slog.Info("applying update", "target_slot", passiveSlot)

	// Check for update
	stageDir := "/tmp/kubesolo-update-stage"
	client := image.NewClient(opts.ServerURL, stageDir)
	defer client.Cleanup()

	// Enable signature verification if public key is configured
	if opts.PubKeyPath != "" {
		client.SetPublicKeyPath(opts.PubKeyPath)
		slog.Info("signature verification enabled", "pubkey", opts.PubKeyPath)
	}

	if err := st.Transition(opts.StatePath, state.PhaseChecking, "", ""); err != nil {
		slog.Warn("state transition failed", "phase", state.PhaseChecking, "error", err)
	}

	meta, err := client.CheckForUpdate()
	if err != nil {
		_ = st.RecordError(opts.StatePath, fmt.Errorf("checking for update: %w", err))
		return fmt.Errorf("checking for update: %w", err)
	}

	slog.Info("update available", "version", meta.Version)

	// Now we know the target version — record it (resets attempt count if it
	// differs from the previous attempt's ToVersion).
	if err := st.Transition(opts.StatePath, state.PhaseDownloading, meta.Version, ""); err != nil {
		slog.Warn("state transition failed", "phase", state.PhaseDownloading, "error", err)
	}

	// Download and verify
	staged, err := client.Download(meta)
	if err != nil {
		_ = st.RecordError(opts.StatePath, fmt.Errorf("downloading update: %w", err))
		return fmt.Errorf("downloading update: %w", err)
	}

	// Mount passive partition
	partInfo, err := partition.GetSlotPartition(passiveSlot)
	if err != nil {
		_ = st.RecordError(opts.StatePath, fmt.Errorf("finding passive partition: %w", err))
		return fmt.Errorf("finding passive partition: %w", err)
	}

	mountPoint := "/tmp/kubesolo-passive-" + passiveSlot
	if err := partition.MountReadWrite(partInfo.Device, mountPoint); err != nil {
		_ = st.RecordError(opts.StatePath, fmt.Errorf("mounting passive partition: %w", err))
		return fmt.Errorf("mounting passive partition: %w", err)
	}
	defer partition.Unmount(mountPoint)

	// Write image to passive partition
	if err := partition.WriteSystemImage(mountPoint, staged.VmlinuzPath, staged.InitramfsPath, staged.Version); err != nil {
		_ = st.RecordError(opts.StatePath, fmt.Errorf("writing system image: %w", err))
		return fmt.Errorf("writing system image: %w", err)
	}

	if err := st.Transition(opts.StatePath, state.PhaseStaged, staged.Version, ""); err != nil {
		slog.Warn("state transition failed", "phase", state.PhaseStaged, "error", err)
	}

	fmt.Printf("Update v%s written to slot %s (%s)\n", staged.Version, passiveSlot, partInfo.Device)
	fmt.Println("Run 'kubesolo-update activate' to boot into the new version")

	return nil
}