// tera-update-tool creates signing keys and deterministic, signed local update packages.
// Private keys are host-only artifacts and are never installed on a target device.
package main

import (
	"archive/tar"
	"compress/gzip"
	"crypto/ed25519"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"
	"time"

	"tera-clone/services/tera-api/internal/update"
)

type privateKeyEnvelope struct {
	FormatVersion int    `json:"format_version"`
	KeyID         string `json:"key_id"`
	Algorithm     string `json:"algorithm"`
	PrivateKey    string `json:"private_key"`
}

type publicKeyEnvelope 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"`
}

func main() {
	if len(os.Args) < 2 {
		usage()
	}
	var err error
	switch os.Args[1] {
	case "keygen":
		err = keygen(os.Args[2:])
	case "pack":
		err = pack(os.Args[2:])
	case "pack-recovery":
		err = packRecovery(os.Args[2:])
	case "inspect":
		err = inspect(os.Args[2:])
	default:
		usage()
	}
	if err != nil {
		fmt.Fprintln(os.Stderr, "tera-update-tool:", err)
		os.Exit(1)
	}
}

func usage() {
	fmt.Fprintln(os.Stderr, "usage: tera-update-tool {keygen|pack|pack-recovery|inspect} [flags]")
	os.Exit(2)
}

func keygen(arguments []string) error {
	flags := flag.NewFlagSet("keygen", flag.ContinueOnError)
	outputDirectory := flags.String("output-dir", "", "directory for the host private key and target public anchor")
	keyID := flags.String("key-id", "", "bounded trust anchor id")
	product := flags.String("product", "tera-core-x48-pro", "allowed product")
	profile := flags.String("profile", "fet3568j-c", "allowed hardware profile")
	if err := flags.Parse(arguments); err != nil {
		return err
	}
	if !safeName(*keyID) || strings.TrimSpace(*outputDirectory) == "" || strings.TrimSpace(*product) == "" || strings.TrimSpace(*profile) == "" {
		return errors.New("output-dir, a safe key-id, product, and profile are required")
	}
	if err := os.MkdirAll(*outputDirectory, 0o700); err != nil {
		return err
	}
	privatePath := filepath.Join(*outputDirectory, *keyID+".private.json")
	publicPath := filepath.Join(*outputDirectory, *keyID+".json")
	if exists(privatePath) || exists(publicPath) {
		return errors.New("refusing to overwrite an existing key")
	}
	publicKey, privateKey, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		return err
	}
	privateValue := privateKeyEnvelope{1, *keyID, "ed25519", base64.StdEncoding.EncodeToString(privateKey)}
	publicValue := publicKeyEnvelope{1, *keyID, "ed25519", base64.StdEncoding.EncodeToString(publicKey), []string{*product}, []string{*profile}}
	if err := writeExclusiveJSON(privatePath, privateValue, 0o600); err != nil {
		return err
	}
	if err := writeExclusiveJSON(publicPath, publicValue, 0o644); err != nil {
		_ = os.Remove(privatePath)
		return err
	}
	fmt.Printf("host private key: %s\ntarget public anchor: %s\n", privatePath, publicPath)
	return nil
}

func pack(arguments []string) error {
	flags := flag.NewFlagSet("pack", flag.ContinueOnError)
	privatePath := flags.String("private-key", "", "host private key JSON")
	outputPath := flags.String("output", "", "new package path")
	payloadPath := flags.String("payload", "", "payload file")
	payloadClass := flags.String("payload-class", "application", "config, application, boot, or full-rootfs")
	packageID := flags.String("package-id", "", "package id")
	version := flags.String("version", "", "package semantic version")
	minimumVersion := flags.String("minimum-current-version", "", "minimum device version")
	product := flags.String("product", "tera-core-x48-pro", "target product")
	profile := flags.String("profile", "fet3568j-c", "target hardware profile")
	createdAtValue := flags.String("created-at", "", "RFC3339 creation time (defaults to now)")
	recovery := flags.String("recovery", "Upload a known-good signed package; running firmware remains unchanged.", "operator recovery guidance")
	if err := flags.Parse(arguments); err != nil {
		return err
	}
	if !safeName(*packageID) || *version == "" || *privatePath == "" || *outputPath == "" || *payloadPath == "" {
		return errors.New("private-key, output, payload, package-id, and version are required")
	}
	if *payloadClass != "config" && *payloadClass != "application" && *payloadClass != "boot" && *payloadClass != "full-rootfs" {
		return errors.New("payload-class is invalid")
	}
	if exists(*outputPath) {
		return errors.New("refusing to overwrite an existing package")
	}
	createdAt := time.Now().UTC()
	if *createdAtValue != "" {
		parsed, err := time.Parse(time.RFC3339, *createdAtValue)
		if err != nil {
			return fmt.Errorf("created-at: %w", err)
		}
		createdAt = parsed.UTC()
	}
	privateKeyValue, err := readPrivateKey(*privatePath)
	if err != nil {
		return err
	}
	payloadInfo, err := safeRegularFile(*payloadPath)
	if err != nil {
		return err
	}
	payloadDigest, err := fileSHA256(*payloadPath)
	if err != nil {
		return err
	}
	payloadArchivePath := "payload/" + filepath.Base(*payloadPath)
	manifest := update.Manifest{
		FormatVersion: 1, PackageID: *packageID, Product: *product, HardwareProfile: *profile,
		Version: *version, MinimumCurrentVersion: *minimumVersion, CreatedAt: createdAt, SigningKeyID: privateKeyValue.KeyID,
		Payloads: []update.Payload{{Name: "primary", Path: payloadArchivePath, Class: *payloadClass, Size: payloadInfo.Size(), SHA256: payloadDigest}},
	}
	manifest.Apply.Capability = "validation-only"
	manifest.Apply.RebootRequired = false
	manifest.Recovery.Capability = "external-service"
	manifest.Recovery.Instructions = strings.TrimSpace(*recovery)
	manifestBytes, err := json.Marshal(manifest)
	if err != nil {
		return err
	}
	privateKeyBytes, _ := base64.StdEncoding.DecodeString(privateKeyValue.PrivateKey)
	signatureBytes, err := json.Marshal(signatureEnvelope{Algorithm: "ed25519", KeyID: privateKeyValue.KeyID, Signature: base64.StdEncoding.EncodeToString(ed25519.Sign(ed25519.PrivateKey(privateKeyBytes), manifestBytes))})
	if err != nil {
		return err
	}
	return writePackage(*outputPath, []packagePayload{{sourcePath: *payloadPath, archivePath: payloadArchivePath, info: payloadInfo}}, manifestBytes, signatureBytes, createdAt, "validation-only")
}

func packRecovery(arguments []string) error {
	flags := flag.NewFlagSet("pack-recovery", flag.ContinueOnError)
	privatePath := flags.String("private-key", "", "host private key JSON")
	outputPath := flags.String("output", "", "new recovery package path")
	rootfsPath := flags.String("rootfs-payload", "", "raw or gzip-compressed rootfs image")
	rootfsEncoding := flags.String("rootfs-encoding", "gzip", "raw or gzip")
	bootPath := flags.String("boot-payload", "", "optional raw signed normal boot FIT")
	packageID := flags.String("package-id", "", "package id")
	version := flags.String("version", "", "package semantic version")
	minimumVersion := flags.String("minimum-current-version", "", "minimum device version")
	product := flags.String("product", "tera-core-x48-pro", "target product")
	profile := flags.String("profile", "fet3568j-c", "target hardware profile")
	createdAtValue := flags.String("created-at", "", "RFC3339 creation time (defaults to now)")
	recovery := flags.String("recovery", "Keep both power inputs connected. If recovery reports a failure, leave the device in recovery and follow the COM5 service procedure.", "operator recovery guidance")
	if err := flags.Parse(arguments); err != nil {
		return err
	}
	if !safeName(*packageID) || *version == "" || *privatePath == "" || *outputPath == "" || *rootfsPath == "" {
		return errors.New("private-key, output, rootfs-payload, package-id, and version are required")
	}
	if *rootfsEncoding != "raw" && *rootfsEncoding != "gzip" {
		return errors.New("rootfs-encoding must be raw or gzip")
	}
	if exists(*outputPath) {
		return errors.New("refusing to overwrite an existing package")
	}
	createdAt := time.Now().UTC()
	if *createdAtValue != "" {
		parsed, err := time.Parse(time.RFC3339, *createdAtValue)
		if err != nil {
			return fmt.Errorf("created-at: %w", err)
		}
		createdAt = parsed.UTC()
	}
	privateKeyValue, err := readPrivateKey(*privatePath)
	if err != nil {
		return err
	}
	rootfsInfo, err := safeRegularFile(*rootfsPath)
	if err != nil {
		return err
	}
	rootfsDigest, err := fileSHA256(*rootfsPath)
	if err != nil {
		return err
	}
	expandedSize, expandedDigest, err := expandedMetadata(*rootfsPath, *rootfsEncoding)
	if err != nil {
		return err
	}
	rootfsArchivePath := "payload/" + filepath.Base(*rootfsPath)
	manifest := update.Manifest{
		FormatVersion: 1, PackageID: *packageID, Product: *product, HardwareProfile: *profile,
		Version: *version, MinimumCurrentVersion: *minimumVersion, CreatedAt: createdAt, SigningKeyID: privateKeyValue.KeyID,
		Payloads: []update.Payload{{
			Name: "rootfs", Path: rootfsArchivePath, Class: "full-rootfs", Size: rootfsInfo.Size(), SHA256: rootfsDigest,
			TargetPartition: "rootfs", Encoding: *rootfsEncoding, ExpandedSize: expandedSize, ExpandedSHA256: expandedDigest,
		}},
	}
	payloads := []packagePayload{{sourcePath: *rootfsPath, archivePath: rootfsArchivePath, info: rootfsInfo}}
	if *bootPath != "" {
		bootInfo, err := safeRegularFile(*bootPath)
		if err != nil {
			return err
		}
		bootDigest, err := fileSHA256(*bootPath)
		if err != nil {
			return err
		}
		bootArchivePath := "payload/" + filepath.Base(*bootPath)
		if bootArchivePath == rootfsArchivePath {
			return errors.New("boot and rootfs payload basenames must differ")
		}
		manifest.Payloads = append(manifest.Payloads, update.Payload{
			Name: "boot", Path: bootArchivePath, Class: "boot", Size: bootInfo.Size(), SHA256: bootDigest,
			TargetPartition: "boot", Encoding: "raw", ExpandedSize: bootInfo.Size(), ExpandedSHA256: bootDigest,
		})
		payloads = append(payloads, packagePayload{sourcePath: *bootPath, archivePath: bootArchivePath, info: bootInfo})
	}
	manifest.Apply.Capability = update.ApplyCapabilityRecoveryV1
	manifest.Apply.RebootRequired = true
	manifest.Recovery.Capability = update.RecoveryCapabilityOfflineV1
	manifest.Recovery.Instructions = strings.TrimSpace(*recovery)
	manifestBytes, err := json.Marshal(manifest)
	if err != nil {
		return err
	}
	privateKeyBytes, _ := base64.StdEncoding.DecodeString(privateKeyValue.PrivateKey)
	signatureBytes, err := json.Marshal(signatureEnvelope{Algorithm: "ed25519", KeyID: privateKeyValue.KeyID, Signature: base64.StdEncoding.EncodeToString(ed25519.Sign(ed25519.PrivateKey(privateKeyBytes), manifestBytes))})
	if err != nil {
		return err
	}
	return writePackage(*outputPath, payloads, manifestBytes, signatureBytes, createdAt, "recovery-v1")
}

func inspect(arguments []string) error {
	flags := flag.NewFlagSet("inspect", flag.ContinueOnError)
	packagePath := flags.String("package", "", "package to verify")
	trustDirectory := flags.String("trust-dir", "", "public trust anchor directory")
	product := flags.String("product", "tera-core-x48-pro", "expected product")
	profile := flags.String("profile", "fet3568j-c", "expected hardware profile")
	currentVersion := flags.String("current-version", "0.1.0-dev", "current device version")
	if err := flags.Parse(arguments); err != nil {
		return err
	}
	manifest, signer, err := update.VerifyPackage(*packagePath, *trustDirectory, *product, *profile, *currentVersion)
	if err != nil {
		return err
	}
	return json.NewEncoder(os.Stdout).Encode(map[string]any{"valid": true, "signer_key_id": signer, "manifest": manifest})
}

func readPrivateKey(path string) (privateKeyEnvelope, error) {
	info, err := safeRegularFile(path)
	if err != nil {
		return privateKeyEnvelope{}, err
	}
	if info.Size() > 4096 {
		return privateKeyEnvelope{}, errors.New("private key file is too large")
	}
	data, err := os.ReadFile(path)
	if err != nil {
		return privateKeyEnvelope{}, err
	}
	var value privateKeyEnvelope
	if err := json.Unmarshal(data, &value); err != nil || value.FormatVersion != 1 || value.Algorithm != "ed25519" || !safeName(value.KeyID) {
		return privateKeyEnvelope{}, errors.New("private key envelope is invalid")
	}
	decoded, err := base64.StdEncoding.DecodeString(value.PrivateKey)
	if err != nil || len(decoded) != ed25519.PrivateKeySize {
		return privateKeyEnvelope{}, errors.New("private key material is invalid")
	}
	return value, nil
}

type packagePayload struct {
	sourcePath  string
	archivePath string
	info        os.FileInfo
}

func writePackage(outputPath string, payloads []packagePayload, manifest, signature []byte, createdAt time.Time, capability string) error {
	parent := filepath.Dir(outputPath)
	if _, err := safeRegularDirectory(parent); err != nil {
		return err
	}
	temporary, err := os.CreateTemp(parent, ".tera-update-*.part")
	if err != nil {
		return err
	}
	temporaryPath := temporary.Name()
	committed := false
	defer func() {
		_ = temporary.Close()
		if !committed {
			_ = os.Remove(temporaryPath)
		}
	}()
	if err := temporary.Chmod(0o600); err != nil {
		return err
	}
	writer := tar.NewWriter(temporary)
	for _, entry := range []struct {
		name string
		data []byte
	}{{"manifest.json", manifest}, {"manifest.sig", signature}} {
		if err := writeTarBytes(writer, entry.name, entry.data, createdAt); err != nil {
			return err
		}
	}
	for _, item := range payloads {
		if err := writer.WriteHeader(&tar.Header{Name: item.archivePath, Mode: 0o600, Size: item.info.Size(), ModTime: createdAt}); err != nil {
			return err
		}
		payload, err := os.Open(item.sourcePath)
		if err != nil {
			return err
		}
		_, copyErr := io.Copy(writer, payload)
		closePayloadErr := payload.Close()
		if copyErr != nil {
			return copyErr
		}
		if closePayloadErr != nil {
			return closePayloadErr
		}
	}
	if err := writer.Close(); err != nil {
		return err
	}
	if err := temporary.Sync(); err != nil {
		return err
	}
	if err := temporary.Close(); err != nil {
		return err
	}
	if err := os.Rename(temporaryPath, outputPath); err != nil {
		return err
	}
	committed = true
	fmt.Printf("signed %s package: %s\n", capability, outputPath)
	return nil
}

func expandedMetadata(payloadPath, encoding string) (int64, string, error) {
	file, err := os.Open(payloadPath)
	if err != nil {
		return 0, "", err
	}
	defer file.Close()
	var reader io.Reader = file
	if encoding == "gzip" {
		compressed, err := gzip.NewReader(file)
		if err != nil {
			return 0, "", fmt.Errorf("open rootfs gzip payload: %w", err)
		}
		defer compressed.Close()
		reader = compressed
	} else if encoding != "raw" {
		return 0, "", fmt.Errorf("unsupported payload encoding %q", encoding)
	}
	hash := sha256.New()
	size, err := io.Copy(hash, reader)
	if err != nil {
		return 0, "", fmt.Errorf("read expanded payload: %w", err)
	}
	if size <= 0 {
		return 0, "", errors.New("expanded payload is empty")
	}
	return size, hex.EncodeToString(hash.Sum(nil)), nil
}

func writeTarBytes(writer *tar.Writer, name string, data []byte, createdAt time.Time) error {
	if err := writer.WriteHeader(&tar.Header{Name: name, Mode: 0o600, Size: int64(len(data)), ModTime: createdAt}); err != nil {
		return err
	}
	_, err := writer.Write(data)
	return err
}

func fileSHA256(path string) (string, error) {
	file, err := os.Open(path)
	if err != nil {
		return "", err
	}
	defer file.Close()
	hash := sha256.New()
	if _, err := io.Copy(hash, file); err != nil {
		return "", err
	}
	return hex.EncodeToString(hash.Sum(nil)), nil
}

func safeRegularFile(path string) (os.FileInfo, error) {
	info, err := os.Lstat(path)
	if err != nil {
		return nil, err
	}
	if !info.Mode().IsRegular() || info.Mode()&os.ModeSymlink != 0 {
		return nil, errors.New("path must be a regular non-symlink file")
	}
	return info, nil
}

func safeRegularDirectory(path string) (os.FileInfo, error) {
	info, err := os.Lstat(path)
	if err != nil {
		return nil, err
	}
	if !info.IsDir() || info.Mode()&os.ModeSymlink != 0 {
		return nil, errors.New("path must be a non-symlink directory")
	}
	return info, nil
}

func writeExclusiveJSON(path string, value any, mode os.FileMode) error {
	file, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_EXCL, mode)
	if err != nil {
		return err
	}
	encodeErr := json.NewEncoder(file).Encode(value)
	syncErr := file.Sync()
	closeErr := file.Close()
	if encodeErr != nil {
		return encodeErr
	}
	if syncErr != nil {
		return syncErr
	}
	return closeErr
}

func exists(path string) bool { _, err := os.Lstat(path); return err == nil }

func safeName(value string) bool {
	if len(value) < 1 || len(value) > 128 {
		return false
	}
	for index, character := range value {
		if character >= 'a' && character <= 'z' || character >= 'A' && character <= 'Z' || character >= '0' && character <= '9' || index > 0 && (character == '.' || character == '_' || character == '-') {
			continue
		}
		return false
	}
	return true
}