// Package update implements the optional self-update: GET /agent/latest, download the asset for GOOS-GOARCH,
// verify its sha256, atomically replace the running binary and ask the caller to exit 0 (the supervisor —
// launchd / systemd — restarts the new binary).
package update

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"time"

	"internetpressure.io/probe-agent/internal/client"
	"internetpressure.io/probe-agent/internal/protocol"
)

// Defaults.
const (
	CheckEvery      = 6 * time.Hour
	FirstCheckAfter = 2 * time.Minute
	DownloadTimeout = 5 * time.Minute
	MaxAssetBytes   = 256 << 20
	lastAttemptFile = "update.last"
	tmpFile         = "update.tmp"
)

// Fetcher is the subset of client.Client used (interface for tests).
type Fetcher interface {
	AgentLatest(ctx context.Context) (*protocol.AgentLatest, error)
	Download(ctx context.Context, url string, w io.Writer, limit int64) (int64, error)
}

// Updater checks for and applies releases.
type Updater struct {
	Fetcher Fetcher
	Version string
	DataDir string
	// ExePath is the binary to replace; "" → os.Executable() (resolved through symlinks).
	ExePath string
	Log     *slog.Logger
}

// Platform key used in the assets map.
func Platform() string { return runtime.GOOS + "-" + runtime.GOARCH }

// Check fetches /agent/latest and applies an update when one is available. Returns true when the binary was
// replaced (the caller must exit 0).
func (u *Updater) Check(ctx context.Context) (bool, error) {
	log := u.Log
	if log == nil {
		log = slog.Default()
	}
	latest, err := u.Fetcher.AgentLatest(ctx)
	if err != nil {
		return false, err
	}
	if latest.Version == "" || latest.Version == u.Version {
		return false, nil
	}
	asset, ok := latest.Assets[Platform()]
	if !ok || asset.URL == "" || asset.SHA256 == "" {
		log.Info("update available but no asset for this platform", "version", latest.Version, "platform", Platform())
		return false, nil
	}
	// Downgrade/loop guard: never attempt the same version twice (covers a bad release that keeps crashing
	// and being "updated" to, or a server offering an older version).
	if last := u.lastAttempt(); last == latest.Version {
		log.Debug("skipping update already attempted", "version", latest.Version)
		return false, nil
	}
	if err := u.recordAttempt(latest.Version); err != nil {
		return false, err
	}
	log.Info("downloading update", "from", u.Version, "to", latest.Version, "url", asset.URL)

	exe := u.ExePath
	if exe == "" {
		p, err := os.Executable()
		if err != nil {
			return false, err
		}
		if r, err := filepath.EvalSymlinks(p); err == nil {
			p = r
		}
		exe = p
	}

	tmp := filepath.Join(u.DataDir, tmpFile)
	if err := u.download(ctx, asset, tmp); err != nil {
		os.Remove(tmp)
		return false, err
	}
	if err := os.Chmod(tmp, 0o755); err != nil {
		os.Remove(tmp)
		return false, err
	}
	// Atomic replace: rename within the same filesystem. If data_dir is on another filesystem, stage a copy
	// next to the binary first.
	if err := os.Rename(tmp, exe); err != nil {
		staged := exe + ".new"
		if cerr := copyFile(tmp, staged, 0o755); cerr != nil {
			os.Remove(tmp)
			return false, fmt.Errorf("stage update: %w", cerr)
		}
		os.Remove(tmp)
		if err := os.Rename(staged, exe); err != nil {
			os.Remove(staged)
			return false, fmt.Errorf("replace binary: %w", err)
		}
	}
	log.Info("update installed, exiting for restart", "version", latest.Version, "binary", exe)
	return true, nil
}

func (u *Updater) download(ctx context.Context, asset protocol.Asset, dst string) error {
	ctx, cancel := context.WithTimeout(ctx, DownloadTimeout)
	defer cancel()
	f, err := os.OpenFile(dst, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0o700)
	if err != nil {
		return err
	}
	h := sha256.New()
	n, err := u.Fetcher.Download(ctx, asset.URL, io.MultiWriter(f, h), MaxAssetBytes)
	if cerr := f.Close(); err == nil {
		err = cerr
	}
	if err != nil {
		return fmt.Errorf("download: %w", err)
	}
	if n == 0 {
		return errors.New("download: empty asset")
	}
	got := hex.EncodeToString(h.Sum(nil))
	if !strings.EqualFold(got, strings.TrimSpace(asset.SHA256)) {
		return fmt.Errorf("download: sha256 mismatch (got %s, want %s)", got, asset.SHA256)
	}
	return nil
}

func (u *Updater) lastAttempt() string {
	b, err := os.ReadFile(filepath.Join(u.DataDir, lastAttemptFile))
	if err != nil {
		return ""
	}
	return strings.TrimSpace(string(b))
}

func (u *Updater) recordAttempt(v string) error {
	return os.WriteFile(filepath.Join(u.DataDir, lastAttemptFile), []byte(v+"\n"), 0o600)
}

func copyFile(src, dst string, mode os.FileMode) error {
	in, err := os.Open(src)
	if err != nil {
		return err
	}
	defer in.Close()
	out, err := os.OpenFile(dst, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, mode)
	if err != nil {
		return err
	}
	if _, err := io.Copy(out, in); err != nil {
		out.Close()
		return err
	}
	return out.Close()
}

// Run checks after FirstCheckAfter and then every CheckEvery; onUpdated is called once the binary was replaced.
func (u *Updater) Run(ctx context.Context, onUpdated func()) {
	log := u.Log
	if log == nil {
		log = slog.Default()
	}
	wait := FirstCheckAfter
	for {
		select {
		case <-ctx.Done():
			return
		case <-time.After(wait):
		}
		wait = CheckEvery
		updated, err := u.Check(ctx)
		if err != nil {
			var he *client.HTTPError
			if errors.As(err, &he) && he.Status == 404 {
				log.Debug("no release feed", "err", err)
			} else {
				log.Warn("update check failed", "err", err)
			}
			continue
		}
		if updated && onUpdated != nil {
			onUpdated()
			return
		}
	}
}