package main

import (
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"syscall"
	"time"

	"forge-mesh/internal/forge"
)

type buildMeta struct {
	ID        string `json:"id"`
	OS        string `json:"os"`
	Arch      string `json:"arch"`
	Version   string `json:"version"`
	Checksum  string `json:"checksum"`
	Signature string `json:"signature"`
}

func (a *Agent) maybeSelfUpdate(pubKeyHex string) {
	if pubKeyHex == "" {
		return
	}
	if err := a.checkAndApplyUpdate(pubKeyHex); err != nil {
		log.Printf("self-update: %v", err)
	}
}

func (a *Agent) checkAndApplyUpdate(pubKeyHex string) error {
	meta, err := a.fetchLatestBuild()
	if err != nil {
		return err
	}
	if meta.Checksum == "" {
		return fmt.Errorf("build metadata missing checksum")
	}

	exe, err := os.Executable()
	if err != nil {
		return err
	}
	exe, err = filepath.EvalSymlinks(exe)
	if err != nil {
		return err
	}

	data, err := os.ReadFile(exe)
	if err != nil {
		return err
	}
	current := sha256.Sum256(data)
	if hex.EncodeToString(current[:]) == meta.Checksum {
		return nil
	}

	log.Printf("self-update: new build %s (%s) available", meta.ID, meta.Version)

	tmp, err := os.CreateTemp(filepath.Dir(exe), "forge-mesh-agent-update-*")
	if err != nil {
		return err
	}
	tmpPath := tmp.Name()
	defer os.Remove(tmpPath)

	downloadURL := fmt.Sprintf("%s/api/v1/public/download/%s", a.serverURL, meta.ID)
	resp, err := http.Get(downloadURL)
	if err != nil {
		tmp.Close()
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		tmp.Close()
		return fmt.Errorf("download status %d", resp.StatusCode)
	}
	if _, err := io.Copy(tmp, resp.Body); err != nil {
		tmp.Close()
		return err
	}
	if err := tmp.Close(); err != nil {
		return err
	}

	artifact, err := os.ReadFile(tmpPath)
	if err != nil {
		return err
	}
	sum := sha256.Sum256(artifact)
	if hex.EncodeToString(sum[:]) != meta.Checksum {
		return fmt.Errorf("downloaded checksum mismatch")
	}

	pub, err := forge.ParsePublicKeyHex(pubKeyHex)
	if err != nil {
		return err
	}
	if meta.Signature != "" && !forge.Verify(pub, artifact, meta.Signature) {
		return fmt.Errorf("signature verification failed")
	}

	if err := os.Chmod(tmpPath, 0o755); err != nil {
		return err
	}
	if err := os.Rename(tmpPath, exe); err != nil {
		return fmt.Errorf("replace binary: %w (restart via systemd)", err)
	}

	log.Printf("self-update: applied build %s, restarting", meta.ID)
	go restartAgent()
	return nil
}

func (a *Agent) fetchLatestBuild() (*buildMeta, error) {
	url := fmt.Sprintf("%s/api/v1/public/builds/latest?os=linux&arch=%s", a.serverURL, runtime.GOARCH)
	resp, err := http.Get(url)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("latest build status %d", resp.StatusCode)
	}
	var meta buildMeta
	if err := json.NewDecoder(resp.Body).Decode(&meta); err != nil {
		return nil, err
	}
	return &meta, nil
}

func restartAgent() {
	time.Sleep(500 * time.Millisecond)
	if os.Getenv("INVOCATION_ID") != "" {
		_ = exec.Command("systemctl", "restart", "forge-mesh-agent").Run()
		return
	}
	exe, err := os.Executable()
	if err != nil {
		os.Exit(0)
	}
	_ = syscall.Exec(exe, append([]string{exe}, os.Args[1:]...), os.Environ())
	os.Exit(0)
}

func verifyAgentUpdate(pubKeyHex string, artifact []byte, sigB64 string) bool {
	pub, err := forge.ParsePublicKeyHex(pubKeyHex)
	if err != nil {
		return false
	}
	return forge.Verify(pub, artifact, sigB64)
}