// Package image handles downloading, verifying, and staging OS update images.
//
// Update images are distributed as pairs of files:
//   - vmlinuz (kernel)
//   - kubesolo-os.gz (initramfs)
//
// These are fetched from an HTTP(S) server that provides a metadata file
// (latest.json) describing available updates.
package image

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

	"github.com/portainer/kubesolo-os/update/pkg/signing"
)

// UpdateMetadata describes an available update from the update server.
type UpdateMetadata struct {
	Version         string `json:"version"`
	VmlinuzURL      string `json:"vmlinuz_url"`
	VmlinuzSHA256   string `json:"vmlinuz_sha256"`
	VmlinuzSigURL   string `json:"vmlinuz_sig_url,omitempty"`
	InitramfsURL    string `json:"initramfs_url"`
	InitramfsSHA256 string `json:"initramfs_sha256"`
	InitramfsSigURL string `json:"initramfs_sig_url,omitempty"`
	MetadataSigURL  string `json:"metadata_sig_url,omitempty"`
	ReleaseNotes    string `json:"release_notes,omitempty"`
	ReleaseDate     string `json:"release_date,omitempty"`
}

// StagedImage represents downloaded and verified update files.
type StagedImage struct {
	VmlinuzPath  string
	InitramfsPath string
	Version       string
}

// Client handles communication with the update server.
type Client struct {
	serverURL    string
	httpClient   *http.Client
	stageDir     string
	pubKeyPath   string // path to Ed25519 public key for signature verification
}

// NewClient creates a new update image client.
func NewClient(serverURL, stageDir string) *Client {
	return &Client{
		serverURL: serverURL,
		httpClient: &http.Client{
			Timeout: 5 * time.Minute,
		},
		stageDir: stageDir,
	}
}

// SetPublicKeyPath sets the path to the Ed25519 public key used
// for verifying update signatures. If set, downloaded images will
// be verified against their .sig files from the update server.
func (c *Client) SetPublicKeyPath(path string) {
	c.pubKeyPath = path
}

// CheckForUpdate fetches the latest update metadata from the server.
func (c *Client) CheckForUpdate() (*UpdateMetadata, error) {
	url := c.serverURL + "/latest.json"
	slog.Info("checking for update", "url", url)

	resp, err := c.httpClient.Get(url)
	if err != nil {
		return nil, fmt.Errorf("fetching update metadata: %w", err)
	}
	defer resp.Body.Close()

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

	var meta UpdateMetadata
	if err := json.NewDecoder(resp.Body).Decode(&meta); err != nil {
		return nil, fmt.Errorf("parsing update metadata: %w", err)
	}

	if meta.Version == "" {
		return nil, fmt.Errorf("update metadata missing version")
	}

	return &meta, nil
}

// Download fetches the update files and verifies their checksums.
func (c *Client) Download(meta *UpdateMetadata) (*StagedImage, error) {
	if err := os.MkdirAll(c.stageDir, 0o755); err != nil {
		return nil, fmt.Errorf("creating stage dir: %w", err)
	}

	vmlinuzPath := filepath.Join(c.stageDir, "vmlinuz")
	initramfsPath := filepath.Join(c.stageDir, "kubesolo-os.gz")

	slog.Info("downloading vmlinuz", "url", meta.VmlinuzURL)
	if err := c.downloadAndVerify(meta.VmlinuzURL, vmlinuzPath, meta.VmlinuzSHA256); err != nil {
		return nil, fmt.Errorf("downloading vmlinuz: %w", err)
	}

	slog.Info("downloading initramfs", "url", meta.InitramfsURL)
	if err := c.downloadAndVerify(meta.InitramfsURL, initramfsPath, meta.InitramfsSHA256); err != nil {
		return nil, fmt.Errorf("downloading initramfs: %w", err)
	}

	// Verify signatures if public key is configured
	if c.pubKeyPath != "" {
		if err := c.verifySignatures(meta, vmlinuzPath, initramfsPath); err != nil {
			os.Remove(vmlinuzPath)
			os.Remove(initramfsPath)
			return nil, fmt.Errorf("signature verification: %w", err)
		}
	}

	return &StagedImage{
		VmlinuzPath:   vmlinuzPath,
		InitramfsPath: initramfsPath,
		Version:       meta.Version,
	}, nil
}

// Cleanup removes staged update files.
func (c *Client) Cleanup() error {
	return os.RemoveAll(c.stageDir)
}

func (c *Client) downloadAndVerify(url, dest, expectedSHA256 string) error {
	resp, err := c.httpClient.Get(url)
	if err != nil {
		return fmt.Errorf("downloading %s: %w", url, err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("server returned %d for %s", resp.StatusCode, url)
	}

	f, err := os.Create(dest)
	if err != nil {
		return fmt.Errorf("creating %s: %w", dest, err)
	}
	defer f.Close()

	hasher := sha256.New()
	writer := io.MultiWriter(f, hasher)

	written, err := io.Copy(writer, resp.Body)
	if err != nil {
		os.Remove(dest)
		return fmt.Errorf("writing %s: %w", dest, err)
	}

	if err := f.Close(); err != nil {
		return fmt.Errorf("closing %s: %w", dest, err)
	}

	// Verify checksum
	if expectedSHA256 != "" {
		actual := hex.EncodeToString(hasher.Sum(nil))
		if actual != expectedSHA256 {
			os.Remove(dest)
			return fmt.Errorf("checksum mismatch for %s: expected %s, got %s", dest, expectedSHA256, actual)
		}
		slog.Debug("checksum verified", "file", dest, "sha256", actual)
	}

	slog.Info("downloaded", "file", dest, "size", written)
	return nil
}

// verifySignatures downloads .sig files and verifies them against the staged images.
func (c *Client) verifySignatures(meta *UpdateMetadata, vmlinuzPath, initramfsPath string) error {
	verifier, err := signing.NewVerifierFromFile(c.pubKeyPath)
	if err != nil {
		return fmt.Errorf("loading public key: %w", err)
	}

	// Verify vmlinuz signature
	if meta.VmlinuzSigURL != "" {
		sigPath := vmlinuzPath + ".sig"
		if err := c.downloadToFile(meta.VmlinuzSigURL, sigPath); err != nil {
			return fmt.Errorf("downloading vmlinuz signature: %w", err)
		}
		if err := verifier.VerifyFile(vmlinuzPath, sigPath); err != nil {
			return fmt.Errorf("vmlinuz: %w", err)
		}
		slog.Info("vmlinuz signature verified")
	}

	// Verify initramfs signature
	if meta.InitramfsSigURL != "" {
		sigPath := initramfsPath + ".sig"
		if err := c.downloadToFile(meta.InitramfsSigURL, sigPath); err != nil {
			return fmt.Errorf("downloading initramfs signature: %w", err)
		}
		if err := verifier.VerifyFile(initramfsPath, sigPath); err != nil {
			return fmt.Errorf("initramfs: %w", err)
		}
		slog.Info("initramfs signature verified")
	}

	return nil
}

// downloadToFile downloads a URL to a local file (used for signature files).
func (c *Client) downloadToFile(url, dest string) error {
	resp, err := c.httpClient.Get(url)
	if err != nil {
		return fmt.Errorf("downloading %s: %w", url, err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("server returned %d for %s", resp.StatusCode, url)
	}

	data, err := io.ReadAll(resp.Body)
	if err != nil {
		return fmt.Errorf("reading response: %w", err)
	}

	return os.WriteFile(dest, data, 0o644)
}

// VerifyFile checks the SHA256 checksum of an existing file.
func VerifyFile(path, expectedSHA256 string) error {
	f, err := os.Open(path)
	if err != nil {
		return err
	}
	defer f.Close()

	hasher := sha256.New()
	if _, err := io.Copy(hasher, f); err != nil {
		return err
	}

	actual := hex.EncodeToString(hasher.Sum(nil))
	if actual != expectedSHA256 {
		return fmt.Errorf("checksum mismatch: expected %s, got %s", expectedSHA256, actual)
	}
	return nil
}