update.go

package main

import (
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"runtime"
	"strconv"
	"strings"
	"time"
)

type githubRelease struct {
	TagName string        `json:"tag_name"`
	Assets  []githubAsset `json:"assets"`
}

type githubAsset struct {
	Name               string `json:"name"`
	BrowserDownloadURL string `json:"browser_download_url"`
}

type semver struct {
	Major int
	Minor int
	Patch int
}

func parseSemver(v string) (semver, error) {
	v = strings.TrimPrefix(v, "v")
	parts := strings.Split(v, ".")
	if len(parts) != 3 {
		return semver{}, fmt.Errorf("invalid version: %s", v)
	}
	major, err := strconv.Atoi(parts[0])
	if err != nil {
		return semver{}, err
	}
	minor, err := strconv.Atoi(parts[1])
	if err != nil {
		return semver{}, err
	}
	patch, err := strconv.Atoi(parts[2])
	if err != nil {
		return semver{}, err
	}
	return semver{major, minor, patch}, nil
}

func (s semver) String() string {
	return fmt.Sprintf("v%d.%d.%d", s.Major, s.Minor, s.Patch)
}

// updateType returns "major", "minor", "patch", or "" if no update needed
func (s semver) updateType(latest semver) string {
	if latest.Major > s.Major {
		return "major"
	}
	if latest.Major < s.Major {
		return ""
	}
	if latest.Minor > s.Minor {
		return "minor"
	}
	if latest.Minor < s.Minor {
		return ""
	}
	if latest.Patch > s.Patch {
		return "patch"
	}
	return ""
}

func getAssetSuffix() string {
	goos := runtime.GOOS
	goarch := runtime.GOARCH

	if goos == "darwin" {
		goos = "macos"
	}

	suffix := goos + "-" + goarch
	if runtime.GOOS == "windows" {
		suffix += ".exe"
	}
	return suffix
}

func checkLatestVersion() (*githubRelease, error) {
	client := &http.Client{Timeout: 3 * time.Second}
	resp, err := client.Get("https://api.github.com/repos/softwarity/aipilot-cli/releases/latest")
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("GitHub API returned %d", resp.StatusCode)
	}

	var release githubRelease
	if err := json.NewDecoder(resp.Body).Decode(&release); err != nil {
		return nil, err
	}
	return &release, nil
}

func findDownloadURL(release *githubRelease) string {
	suffix := getAssetSuffix()
	for _, asset := range release.Assets {
		if strings.HasSuffix(asset.Name, suffix) {
			return asset.BrowserDownloadURL
		}
	}
	return ""
}

func getExecutablePath() (string, error) {
	exe, err := os.Executable()
	if err != nil {
		return "", err
	}
	return filepath.EvalSymlinks(exe)
}

func findChecksumsURL(release *githubRelease) string {
	for _, asset := range release.Assets {
		if asset.Name == "checksums.txt" {
			return asset.BrowserDownloadURL
		}
	}
	return ""
}

func verifyChecksum(filePath string, release *githubRelease) error {
	checksumsURL := findChecksumsURL(release)
	if checksumsURL == "" {
		fmt.Printf("%s  Warning: No checksums.txt in release, skipping verification%s\n", yellow, reset)
		return nil
	}

	client := &http.Client{Timeout: 10 * time.Second}
	resp, err := client.Get(checksumsURL)
	if err != nil {
		return fmt.Errorf("failed to download checksums: %w", err)
	}
	defer resp.Body.Close()

	checksumsData, err := io.ReadAll(resp.Body)
	if err != nil {
		return fmt.Errorf("failed to read checksums: %w", err)
	}

	// Parse checksums.txt (format: "sha256  filename")
	suffix := getAssetSuffix()
	var expectedHash string
	for _, line := range strings.Split(string(checksumsData), "\n") {
		parts := strings.Fields(line)
		if len(parts) == 2 && strings.HasSuffix(parts[1], suffix) {
			expectedHash = parts[0]
			break
		}
	}
	if expectedHash == "" {
		return fmt.Errorf("no checksum found for %s in checksums.txt", suffix)
	}

	f, err := os.Open(filePath)
	if err != nil {
		return fmt.Errorf("cannot open file for checksum: %w", err)
	}
	defer f.Close()

	h := sha256.New()
	if _, err := io.Copy(h, f); err != nil {
		return fmt.Errorf("failed to compute checksum: %w", err)
	}
	actualHash := hex.EncodeToString(h.Sum(nil))

	if actualHash != expectedHash {
		return fmt.Errorf("checksum mismatch: expected %s, got %s", expectedHash, actualHash)
	}
	return nil
}

func downloadAndReplace(downloadURL, exePath string, release *githubRelease) error {
	resp, err := http.Get(downloadURL)
	if err != nil {
		return fmt.Errorf("download failed: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("download failed: HTTP %d", resp.StatusCode)
	}

	// Write to temp file in same directory (ensures same filesystem for rename)
	dir := filepath.Dir(exePath)
	tmpFile, err := os.CreateTemp(dir, ".aipilot-cli-update-*")
	if err != nil {
		return fmt.Errorf("cannot create temp file: %w", err)
	}
	tmpPath := tmpFile.Name()

	if _, err := io.Copy(tmpFile, resp.Body); err != nil {
		tmpFile.Close()
		os.Remove(tmpPath)
		return fmt.Errorf("download interrupted: %w", err)
	}
	tmpFile.Close()

	if err := os.Chmod(tmpPath, 0755); err != nil {
		os.Remove(tmpPath)
		return err
	}

	// Verify checksum before replacing
	if release != nil {
		if err := verifyChecksum(tmpPath, release); err != nil {
			os.Remove(tmpPath)
			return fmt.Errorf("checksum verification failed: %w", err)
		}
	}

	// Replace the binary
	if runtime.GOOS == "windows" {
		// Windows: can't delete running binary, but can rename it
		oldPath := exePath + ".old"
		os.Remove(oldPath)
		if err := os.Rename(exePath, oldPath); err != nil {
			os.Remove(tmpPath)
			return fmt.Errorf("cannot rename current binary: %w", err)
		}
		if err := os.Rename(tmpPath, exePath); err != nil {
			os.Rename(oldPath, exePath) // restore
			os.Remove(tmpPath)
			return fmt.Errorf("cannot install new binary: %w", err)
		}
	} else {
		// Unix: can replace running binary directly
		if err := os.Rename(tmpPath, exePath); err != nil {
			os.Remove(tmpPath)
			return fmt.Errorf("cannot replace binary: %w", err)
		}
	}

	return nil
}

// cleanupOldBinary removes leftover .old file from Windows update
func cleanupOldBinary() {
	if runtime.GOOS != "windows" {
		return
	}
	if exe, err := getExecutablePath(); err == nil {
		os.Remove(exe + ".old")
	}
}

// checkUpdateOnStartup checks for updates at startup and prompts the user.
func checkUpdateOnStartup() {
	current, err := parseSemver(Version)
	if err != nil {
		return // dev build, skip
	}

	fmt.Printf("%sChecking for updates...%s\r", dim, reset)

	release, err := checkLatestVersion()
	if err != nil {
		fmt.Printf("                       \r") // clear line
		return
	}

	latest, err := parseSemver(release.TagName)
	if err != nil {
		fmt.Printf("                       \r")
		return
	}

	updateType := current.updateType(latest)
	if updateType == "" {
		fmt.Printf("                       \r")
		return
	}

	downloadURL := findDownloadURL(release)
	if downloadURL == "" {
		return
	}

	exePath, err := getExecutablePath()
	if err != nil {
		return
	}

	fmt.Printf("%s⬆ Update available: %s → %s (%s)%s\n", cyan, current.String(), latest.String(), updateType, reset)
	fmt.Printf("  Update now? [Y/n] ")

	var answer string
	fmt.Scanln(&answer)
	answer = strings.TrimSpace(strings.ToLower(answer))

	if answer == "n" || answer == "no" {
		fmt.Printf("%s  Skipped.%s\n", dim, reset)
		return
	}

	fmt.Printf("%s  Updating...%s\n", cyan, reset)
	if err := downloadAndReplace(downloadURL, exePath, release); err != nil {
		fmt.Printf("%s  Update failed: %v%s\n", yellow, err, reset)
		return
	}
	fmt.Printf("%s  ✓ Updated to %s. Restarting...%s\n", green, latest.String(), reset)
	restartSelf(exePath)
}

// forceUpdate performs a blocking update check and install (--update flag)
func forceUpdate() {
	current, err := parseSemver(Version)
	if err != nil {
		fmt.Printf("%sCannot check updates: invalid version %q%s\n", yellow, Version, reset)
		return
	}

	fmt.Printf("Current version: %s\n", current.String())
	fmt.Printf("Checking for updates...\n")

	release, err := checkLatestVersion()
	if err != nil {
		fmt.Printf("%sFailed to check: %v%s\n", red, err, reset)
		return
	}

	latest, err := parseSemver(release.TagName)
	if err != nil {
		fmt.Printf("%sInvalid remote version: %s%s\n", red, release.TagName, reset)
		return
	}

	updateType := current.updateType(latest)
	if updateType == "" {
		fmt.Printf("%s✓ Already up to date (%s)%s\n", green, current.String(), reset)
		return
	}

	downloadURL := findDownloadURL(release)
	if downloadURL == "" {
		fmt.Printf("%sNo binary for %s/%s%s\n", yellow, runtime.GOOS, runtime.GOARCH, reset)
		return
	}

	exePath, err := getExecutablePath()
	if err != nil {
		fmt.Printf("%sCannot determine executable path: %v%s\n", red, err, reset)
		return
	}

	fmt.Printf("Updating %s → %s...\n", current.String(), latest.String())
	if err := downloadAndReplace(downloadURL, exePath, release); err != nil {
		fmt.Printf("%sFailed to update: %v%s\n", red, err, reset)
		return
	}
	fmt.Printf("%s✓ Updated to %s%s\n", green, latest.String(), reset)
}