package update

import (
	"archive/tar"
	"bytes"
	"compress/gzip"
	"crypto/ed25519"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"path"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
)

const (
	manifestPath  = "manifest.json"
	signaturePath = "manifest.sig"
	maxEntries    = 32
	maxManifest   = 64 * 1024
	maxSignature  = 4 * 1024
)

var safeIdentifier = regexp.MustCompile(`^[A-Za-z0-9][A-Za-z0-9._-]{0,127}$`)

type trustAnchor struct {
	FormatVersion   int      `json:"format_version"`
	KeyID           string   `json:"key_id"`
	Algorithm       string   `json:"algorithm"`
	PublicKey       string   `json:"public_key"`
	AllowedProducts []string `json:"allowed_products"`
	AllowedProfiles []string `json:"allowed_profiles"`
}

type signatureEnvelope struct {
	Algorithm string `json:"algorithm"`
	KeyID     string `json:"key_id"`
	Signature string `json:"signature"`
}

type scannedEntry struct {
	Size   int64
	SHA256 string
}

func VerifyPackage(packagePath, trustDirectory, product, profile, currentVersion string) (Manifest, string, error) {
	entries, manifestBytes, signatureBytes, err := scanPackage(packagePath)
	if err != nil {
		return Manifest{}, "", err
	}
	var manifest Manifest
	if err := decodeSingleJSON(manifestBytes, &manifest); err != nil {
		return Manifest{}, "", fmt.Errorf("manifest is invalid: %w", err)
	}
	var envelope signatureEnvelope
	if err := decodeSingleJSON(signatureBytes, &envelope); err != nil {
		return Manifest{}, "", fmt.Errorf("signature envelope is invalid: %w", err)
	}
	if err := validateManifest(manifest, entries, product, profile, currentVersion); err != nil {
		return Manifest{}, "", err
	}
	if envelope.Algorithm != "ed25519" || envelope.KeyID == "" || envelope.KeyID != manifest.SigningKeyID {
		return Manifest{}, "", errors.New("package signer metadata does not match the manifest")
	}
	anchor, publicKey, err := loadTrustAnchor(trustDirectory, envelope.KeyID, product, profile)
	if err != nil {
		return Manifest{}, "", err
	}
	signature, err := base64.StdEncoding.DecodeString(envelope.Signature)
	if err != nil || len(signature) != ed25519.SignatureSize || !ed25519.Verify(publicKey, manifestBytes, signature) {
		return Manifest{}, "", errors.New("package signature verification failed")
	}
	return manifest, anchor.KeyID, nil
}

func scanPackage(packagePath string) (map[string]scannedEntry, []byte, []byte, error) {
	file, err := os.Open(packagePath)
	if err != nil {
		return nil, nil, nil, fmt.Errorf("open staged package: %w", err)
	}
	defer file.Close()
	reader := tar.NewReader(file)
	entries := make(map[string]scannedEntry)
	var manifestBytes, signatureBytes []byte
	for count := 0; ; count++ {
		header, err := reader.Next()
		if errors.Is(err, io.EOF) {
			break
		}
		if err != nil {
			return nil, nil, nil, fmt.Errorf("read package archive: %w", err)
		}
		if count >= maxEntries {
			return nil, nil, nil, errors.New("package contains too many entries")
		}
		name := path.Clean(strings.ReplaceAll(header.Name, "\\", "/"))
		if header.Typeflag != tar.TypeReg || name != header.Name || name == "." || strings.HasPrefix(name, "/") || strings.HasPrefix(name, "../") || strings.Contains(name, "/../") {
			return nil, nil, nil, fmt.Errorf("unsafe package entry %q", header.Name)
		}
		if _, exists := entries[name]; exists {
			return nil, nil, nil, fmt.Errorf("duplicate package entry %q", name)
		}
		if header.Size < 0 {
			return nil, nil, nil, fmt.Errorf("invalid package entry size for %q", name)
		}
		limit := header.Size
		if name == manifestPath && limit > maxManifest || name == signaturePath && limit > maxSignature {
			return nil, nil, nil, fmt.Errorf("package metadata entry %q is too large", name)
		}
		hash := sha256.New()
		var metadata bytes.Buffer
		writer := io.Writer(hash)
		if name == manifestPath || name == signaturePath {
			writer = io.MultiWriter(hash, &metadata)
		}
		written, err := io.CopyN(writer, reader, limit)
		if err != nil || written != limit {
			return nil, nil, nil, fmt.Errorf("read package entry %q: %w", name, err)
		}
		entries[name] = scannedEntry{Size: header.Size, SHA256: hex.EncodeToString(hash.Sum(nil))}
		if name == manifestPath {
			manifestBytes = append([]byte(nil), metadata.Bytes()...)
		} else if name == signaturePath {
			signatureBytes = append([]byte(nil), metadata.Bytes()...)
		}
	}
	if len(manifestBytes) == 0 || len(signatureBytes) == 0 {
		return nil, nil, nil, errors.New("package must contain manifest.json and manifest.sig")
	}
	return entries, manifestBytes, signatureBytes, nil
}

func validateManifest(manifest Manifest, entries map[string]scannedEntry, product, profile, currentVersion string) error {
	if manifest.FormatVersion != 1 || !safeIdentifier.MatchString(manifest.PackageID) || !safeIdentifier.MatchString(manifest.SigningKeyID) {
		return errors.New("manifest identity is invalid")
	}
	if manifest.Product != product || manifest.HardwareProfile != profile {
		return fmt.Errorf("package compatibility mismatch: product=%q profile=%q", manifest.Product, manifest.HardwareProfile)
	}
	if manifest.CreatedAt.IsZero() || len(manifest.Payloads) == 0 || len(manifest.Payloads) > 16 {
		return errors.New("manifest payload metadata is invalid")
	}
	if manifest.Recovery.Capability == "" || strings.TrimSpace(manifest.Recovery.Instructions) == "" || len(manifest.Recovery.Instructions) > 512 {
		return errors.New("manifest recovery guidance is required")
	}
	switch manifest.Apply.Capability {
	case ApplyCapabilityValidationOnly:
		if manifest.Apply.RebootRequired {
			return errors.New("validation-only package must not request a reboot")
		}
	case ApplyCapabilityRecoveryV1:
		if !manifest.Apply.RebootRequired || manifest.Recovery.Capability != RecoveryCapabilityOfflineV1 {
			return errors.New("recovery-v1 package must request reboot and declare offline-resumable-v1 recovery")
		}
	default:
		return fmt.Errorf("unsupported apply capability %q", manifest.Apply.Capability)
	}
	packageVersion, err := parseVersion(manifest.Version)
	if err != nil {
		return fmt.Errorf("package version is invalid: %w", err)
	}
	current, err := parseVersion(currentVersion)
	if err != nil {
		return fmt.Errorf("device version is invalid: %w", err)
	}
	if !manifest.AllowDowngrade && compareVersion(packageVersion, current) < 0 {
		return errors.New("package downgrade is not allowed")
	}
	if manifest.MinimumCurrentVersion != "" {
		minimum, err := parseVersion(manifest.MinimumCurrentVersion)
		if err != nil || compareVersion(current, minimum) < 0 {
			return errors.New("device version is below the package minimum")
		}
	}
	wanted := make(map[string]struct{}, len(manifest.Payloads)+2)
	wanted[manifestPath] = struct{}{}
	wanted[signaturePath] = struct{}{}
	recoveryTargets := make(map[string]struct{})
	for _, payload := range manifest.Payloads {
		if !safeIdentifier.MatchString(payload.Name) || !strings.HasPrefix(payload.Path, "payload/") || path.Clean(payload.Path) != payload.Path {
			return errors.New("manifest payload path is invalid")
		}
		if payload.Class != "config" && payload.Class != "application" && payload.Class != "boot" && payload.Class != "full-rootfs" {
			return fmt.Errorf("unsupported payload class %q", payload.Class)
		}
		if payload.Size <= 0 {
			return fmt.Errorf("payload %q has an invalid size", payload.Path)
		}
		if manifest.Apply.Capability == ApplyCapabilityRecoveryV1 {
			if err := validateRecoveryPayload(payload, recoveryTargets); err != nil {
				return err
			}
		} else if payload.TargetPartition != "" || payload.Encoding != "" || payload.ExpandedSize != 0 || payload.ExpandedSHA256 != "" {
			return fmt.Errorf("validation-only payload %q must not declare recovery write metadata", payload.Path)
		}
		if _, duplicate := wanted[payload.Path]; duplicate {
			return fmt.Errorf("duplicate manifest payload %q", payload.Path)
		}
		entry, found := entries[payload.Path]
		if !found || entry.Size != payload.Size || !strings.EqualFold(entry.SHA256, payload.SHA256) {
			return fmt.Errorf("payload digest or size mismatch for %q", payload.Path)
		}
		wanted[payload.Path] = struct{}{}
	}
	if manifest.Apply.Capability == ApplyCapabilityRecoveryV1 {
		if _, found := recoveryTargets["rootfs"]; !found {
			return errors.New("recovery-v1 package must contain one rootfs payload")
		}
		if len(recoveryTargets) != len(manifest.Payloads) || len(recoveryTargets) > 2 {
			return errors.New("recovery-v1 package may contain only one rootfs and one optional boot payload")
		}
	}
	if len(entries) != len(wanted) {
		return errors.New("package contains unlisted entries")
	}
	return nil
}

func validateRecoveryPayload(payload Payload, targets map[string]struct{}) error {
	wantedTarget := ""
	switch payload.Class {
	case "full-rootfs":
		wantedTarget = "rootfs"
	case "boot":
		wantedTarget = "boot"
	default:
		return fmt.Errorf("recovery-v1 payload class %q is not allowlisted", payload.Class)
	}
	if payload.TargetPartition != wantedTarget {
		return fmt.Errorf("payload %q must target partition %q", payload.Path, wantedTarget)
	}
	if _, duplicate := targets[wantedTarget]; duplicate {
		return fmt.Errorf("duplicate recovery target partition %q", wantedTarget)
	}
	if payload.Encoding != "raw" && payload.Encoding != "gzip" {
		return fmt.Errorf("payload %q has unsupported recovery encoding %q", payload.Path, payload.Encoding)
	}
	if payload.ExpandedSize <= 0 || !validSHA256(payload.ExpandedSHA256) {
		return fmt.Errorf("payload %q has invalid expanded size or SHA-256", payload.Path)
	}
	if payload.Encoding == "raw" && (payload.ExpandedSize != payload.Size || !strings.EqualFold(payload.ExpandedSHA256, payload.SHA256)) {
		return fmt.Errorf("raw payload %q expanded metadata must match its archive metadata", payload.Path)
	}
	targets[wantedTarget] = struct{}{}
	return nil
}

func validSHA256(value string) bool {
	if len(value) != sha256.Size*2 {
		return false
	}
	_, err := hex.DecodeString(value)
	return err == nil
}

// StreamPayload rechecks both the archived and expanded payload digests while
// streaming exactly one manifest-approved payload to writer. It never extracts
// a package entry to the recovery filesystem.
func StreamPayload(packagePath string, payload Payload, writer io.Writer) error {
	file, err := os.Open(packagePath)
	if err != nil {
		return fmt.Errorf("open validated package payload: %w", err)
	}
	defer file.Close()
	reader := tar.NewReader(file)
	for {
		header, err := reader.Next()
		if errors.Is(err, io.EOF) {
			return fmt.Errorf("package payload %q is missing", payload.Path)
		}
		if err != nil {
			return fmt.Errorf("read validated package payload: %w", err)
		}
		if header.Name != payload.Path {
			continue
		}
		if header.Typeflag != tar.TypeReg || header.Size != payload.Size {
			return fmt.Errorf("package payload %q type or size changed", payload.Path)
		}
		archivedHash := sha256.New()
		limited := &io.LimitedReader{R: reader, N: payload.Size}
		archivedReader := io.TeeReader(limited, archivedHash)
		expandedHash := sha256.New()
		counter := &countingWriter{writer: io.MultiWriter(writer, expandedHash)}

		var copyErr error
		if payload.Encoding == "gzip" {
			compressed, err := gzip.NewReader(archivedReader)
			if err != nil {
				return fmt.Errorf("open gzip payload %q: %w", payload.Path, err)
			}
			_, copyErr = io.CopyN(counter, compressed, payload.ExpandedSize+1)
			if errors.Is(copyErr, io.EOF) && counter.written == payload.ExpandedSize {
				copyErr = nil
			}
			closeErr := compressed.Close()
			if copyErr == nil && closeErr != nil {
				copyErr = closeErr
			}
			if _, err := io.Copy(io.Discard, archivedReader); copyErr == nil && err != nil {
				copyErr = err
			}
		} else if payload.Encoding == "raw" {
			_, copyErr = io.Copy(counter, archivedReader)
		} else {
			return fmt.Errorf("unsupported payload encoding %q", payload.Encoding)
		}
		if copyErr != nil {
			return fmt.Errorf("stream payload %q: %w", payload.Path, copyErr)
		}
		if limited.N != 0 || counter.written != payload.ExpandedSize ||
			!strings.EqualFold(hex.EncodeToString(archivedHash.Sum(nil)), payload.SHA256) ||
			!strings.EqualFold(hex.EncodeToString(expandedHash.Sum(nil)), payload.ExpandedSHA256) {
			return fmt.Errorf("payload %q changed during recovery streaming", payload.Path)
		}
		return nil
	}
}

type countingWriter struct {
	writer  io.Writer
	written int64
}

func (w *countingWriter) Write(value []byte) (int, error) {
	n, err := w.writer.Write(value)
	w.written += int64(n)
	return n, err
}

func loadTrustAnchor(directory, keyID, product, profile string) (trustAnchor, ed25519.PublicKey, error) {
	if !safeIdentifier.MatchString(keyID) {
		return trustAnchor{}, nil, errors.New("invalid signer key id")
	}
	anchorPath := filepath.Join(directory, keyID+".json")
	info, err := os.Lstat(anchorPath)
	if err != nil || !info.Mode().IsRegular() || info.Mode()&os.ModeSymlink != 0 {
		return trustAnchor{}, nil, errors.New("trusted signer is not installed")
	}
	data, err := os.ReadFile(anchorPath)
	if err != nil {
		return trustAnchor{}, nil, errors.New("trusted signer cannot be read")
	}
	var anchor trustAnchor
	if err := decodeSingleJSON(data, &anchor); err != nil || anchor.FormatVersion != 1 || anchor.KeyID != keyID || anchor.Algorithm != "ed25519" || !contains(anchor.AllowedProducts, product) || !contains(anchor.AllowedProfiles, profile) {
		return trustAnchor{}, nil, errors.New("trusted signer policy is invalid for this device")
	}
	publicKey, err := base64.StdEncoding.DecodeString(anchor.PublicKey)
	if err != nil || len(publicKey) != ed25519.PublicKeySize {
		return trustAnchor{}, nil, errors.New("trusted signer public key is invalid")
	}
	return anchor, ed25519.PublicKey(publicKey), nil
}

func decodeSingleJSON(data []byte, target any) error {
	decoder := json.NewDecoder(bytes.NewReader(data))
	decoder.DisallowUnknownFields()
	if err := decoder.Decode(target); err != nil {
		return err
	}
	var extra any
	if err := decoder.Decode(&extra); !errors.Is(err, io.EOF) {
		return errors.New("must contain one JSON value")
	}
	return nil
}

func contains(values []string, wanted string) bool {
	for _, value := range values {
		if value == wanted {
			return true
		}
	}
	return false
}

type semanticVersion struct {
	major, minor, patch int
	prerelease          string
}

func parseVersion(value string) (semanticVersion, error) {
	core := strings.SplitN(value, "+", 2)[0]
	parts := strings.SplitN(core, "-", 2)
	numbers := strings.Split(parts[0], ".")
	if len(numbers) != 3 {
		return semanticVersion{}, errors.New("expected MAJOR.MINOR.PATCH")
	}
	parsed := semanticVersion{}
	destinations := []*int{&parsed.major, &parsed.minor, &parsed.patch}
	for index, number := range numbers {
		if number == "" || len(number) > 1 && number[0] == '0' {
			return semanticVersion{}, errors.New("invalid numeric component")
		}
		value, err := strconv.Atoi(number)
		if err != nil || value < 0 {
			return semanticVersion{}, errors.New("invalid numeric component")
		}
		*destinations[index] = value
	}
	if len(parts) == 2 {
		if parts[1] == "" || strings.ContainsAny(parts[1], " /\\") {
			return semanticVersion{}, errors.New("invalid prerelease")
		}
		parsed.prerelease = parts[1]
	}
	return parsed, nil
}

func compareVersion(left, right semanticVersion) int {
	for _, pair := range [][2]int{{left.major, right.major}, {left.minor, right.minor}, {left.patch, right.patch}} {
		if pair[0] < pair[1] {
			return -1
		}
		if pair[0] > pair[1] {
			return 1
		}
	}
	if left.prerelease == right.prerelease {
		return 0
	}
	if left.prerelease == "" {
		return 1
	}
	if right.prerelease == "" {
		return -1
	}
	return strings.Compare(left.prerelease, right.prerelease)
}