package update

import (
	"archive/tar"
	"bytes"
	"context"
	"crypto/ed25519"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"errors"
	"os"
	"path/filepath"
	"testing"
	"time"
)

const (
	testProduct = "tera-core-x48-pro"
	testProfile = "fet3568j-c"
	testKeyID   = "test-lab-key"
)

func TestManagerStagesAndValidatesSignedPackage(t *testing.T) {
	manager, privateKey := newTestManager(t)
	packageBytes := signedPackage(t, privateKey, testProfile, []byte("application payload"))
	status, err := manager.Upload(context.Background(), UploadRequest{Filename: "release.bin", ContentLength: int64(len(packageBytes)), Actor: "admin", Body: bytes.NewReader(packageBytes)})
	if err != nil {
		t.Fatalf("Upload: %v", err)
	}
	if status.State != StateValidated || status.Manifest == nil || status.Manifest.Version != "1.1.0" || status.SignerKeyID != testKeyID {
		t.Fatalf("unexpected validated status: %+v", status)
	}
	if status.ApplyEnabled || status.ApplyDisabledReason == "" || status.RecoveryInstructions == "" {
		t.Fatalf("unsafe apply state: %+v", status)
	}
	if status.Capacity.UpdateReserved != 0 {
		t.Fatalf("completed upload retained capacity reservation: %+v", status.Capacity)
	}
	restarted, err := NewManager(testManagerOptions(filepath.Dir(manager.root)))
	if err != nil {
		t.Fatalf("restart manager: %v", err)
	}
	recovered, err := restarted.Status(context.Background())
	if err != nil || recovered.State != StateValidated || recovered.Manifest == nil {
		t.Fatalf("recovered status = %+v, %v", recovered, err)
	}
}

func TestManagerRejectsTamperedAndWrongProfilePackages(t *testing.T) {
	for _, test := range []struct {
		name, profile string
		tamper        bool
	}{
		{name: "tampered", profile: testProfile, tamper: true},
		{name: "wrong-profile", profile: "other-board"},
	} {
		t.Run(test.name, func(t *testing.T) {
			manager, privateKey := newTestManager(t)
			payload := []byte("PAYLOAD-TAMPER-ME")
			packageBytes := signedPackage(t, privateKey, test.profile, payload)
			if test.tamper {
				index := bytes.LastIndex(packageBytes, payload)
				if index < 0 {
					t.Fatal("payload not found in test package")
				}
				packageBytes[index] ^= 0xff
			}
			status, err := manager.Upload(context.Background(), UploadRequest{Filename: "release.bin", ContentLength: int64(len(packageBytes)), Body: bytes.NewReader(packageBytes)})
			if err == nil || Code(err) != "update_validation_failed" || status.State != StateFailed {
				t.Fatalf("status = %+v, err = %v", status, err)
			}
			if status.RecoveryInstructions == "" {
				t.Fatal("failure omitted recovery instructions")
			}
			if status.Capacity.UpdateReserved != 0 {
				t.Fatalf("failed upload retained capacity reservation: %+v", status.Capacity)
			}
		})
	}
}

func TestManagerRecoversInterruptedJournal(t *testing.T) {
	manager, _ := newTestManager(t)
	manager.status.State = StateReceiving
	manager.status.OperationID = "fw-interrupted"
	if err := manager.persistLocked(); err != nil {
		t.Fatal(err)
	}
	if err := os.WriteFile(filepath.Join(manager.incoming, "fw-interrupted.part"), []byte("partial"), 0o600); err != nil {
		t.Fatal(err)
	}
	restarted, err := NewManager(testManagerOptions(filepath.Dir(manager.root)))
	if err != nil {
		t.Fatal(err)
	}
	status, err := restarted.Status(context.Background())
	if err != nil || status.State != StateFailed || status.FailureCode != "update_interrupted" {
		t.Fatalf("status = %+v, err = %v", status, err)
	}
	if _, err := os.Stat(filepath.Join(manager.incoming, "fw-interrupted.part")); !os.IsNotExist(err) {
		t.Fatalf("partial package still exists: %v", err)
	}
}

func TestManagerRecoveryApplyIsFeatureGatedAndArmsRebootOnce(t *testing.T) {
	manager, privateKey := newTestManager(t)
	packageBytes := signedRecoveryPackage(t, privateKey, []byte("rootfs image"))
	status, err := manager.Upload(context.Background(), UploadRequest{Filename: "release.tera-update", ContentLength: int64(len(packageBytes)), Actor: "admin", Body: bytes.NewReader(packageBytes)})
	if err != nil {
		t.Fatal(err)
	}
	if status.ApplyEnabled || status.ApplyDisabledReason == "" {
		t.Fatalf("recovery apply escaped the default gate: %+v", status)
	}
	if _, err := manager.Apply(context.Background(), ApplyRequest{Actor: "admin", OperationID: status.OperationID}); Code(err) != "update_apply_unsupported" {
		t.Fatalf("disabled Apply error = %v", err)
	}

	requested, rebooted := 0, 0
	manager.recoveryApplyEnabled = true
	manager.requestRecovery = func(operationID string) error {
		requested++
		if operationID != status.OperationID {
			t.Fatalf("request operation = %q", operationID)
		}
		return nil
	}
	manager.reboot = func() error { rebooted++; return nil }
	manager.status.ApplyEnabled, manager.status.ApplyDisabledReason = manager.applyPolicy(*manager.status.Manifest)
	if _, err := manager.Apply(context.Background(), ApplyRequest{Actor: "admin", OperationID: "fw-wrong"}); Code(err) != "update_validation_failed" || requested != 0 || rebooted != 0 {
		t.Fatalf("mismatched operation error=%v requested=%d rebooted=%d", err, requested, rebooted)
	}
	result, err := manager.Apply(context.Background(), ApplyRequest{Actor: "admin", OperationID: status.OperationID})
	if err != nil || result.State != StateRebootPending || requested != 1 || rebooted != 1 {
		t.Fatalf("Apply result=%+v error=%v requested=%d rebooted=%d", result, err, requested, rebooted)
	}
	if _, err := manager.Apply(context.Background(), ApplyRequest{Actor: "admin", OperationID: status.OperationID}); Code(err) != "update_apply_unsupported" || requested != 1 || rebooted != 1 {
		t.Fatalf("repeated Apply error=%v requested=%d rebooted=%d", err, requested, rebooted)
	}
}

func TestManagerDoesNotRebootWhenRecoveryRequestFails(t *testing.T) {
	manager, privateKey := newTestManager(t)
	manager.recoveryApplyEnabled = true
	manager.requestRecovery = func(string) error { return errors.New("misc readback mismatch") }
	rebooted := 0
	manager.reboot = func() error { rebooted++; return nil }
	packageBytes := signedRecoveryPackage(t, privateKey, []byte("rootfs image"))
	uploaded, err := manager.Upload(context.Background(), UploadRequest{Filename: "release.tera-update", ContentLength: int64(len(packageBytes)), Actor: "admin", Body: bytes.NewReader(packageBytes)})
	if err != nil {
		t.Fatal(err)
	}
	status, err := manager.Apply(context.Background(), ApplyRequest{Actor: "admin", OperationID: uploaded.OperationID})
	if Code(err) != "update_recovery_request_failed" || status.State != StateFailed || rebooted != 0 {
		t.Fatalf("status=%+v error=%v rebooted=%d", status, err, rebooted)
	}
}

func newTestManager(t *testing.T) (*Manager, ed25519.PrivateKey) {
	t.Helper()
	base := t.TempDir()
	options := testManagerOptions(base)
	for _, directory := range []string{options.UpdateRoot, filepath.Join(options.UpdateRoot, "incoming"), filepath.Join(options.UpdateRoot, "validated"), filepath.Join(options.UpdateRoot, "journal"), options.TrustDirectory} {
		if err := os.MkdirAll(directory, 0o700); err != nil {
			t.Fatal(err)
		}
	}
	policy := map[string]any{
		"format_version": 1, "layout": "tera-userdata-v1",
		"mount":       map[string]any{"source": "/dev/mmcblk0p8", "partlabel": "userdata", "filesystem": "ext4", "required_options": []string{"rw"}},
		"caps_bytes":  map[string]int64{"update_staging": 64 << 20, "update_package_max": 32 << 20, "port_events": 1, "audit_logs": 1, "system_events": 1, "state_and_indexes": 1, "export_scratch": 1, "global_free_floor": 1},
		"enforcement": "application-governor-v1", "persistent_uart_payload": false,
	}
	writeJSON(t, options.PolicyPath, policy)
	publicKey, privateKey, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		t.Fatal(err)
	}
	anchor := map[string]any{"format_version": 1, "key_id": testKeyID, "algorithm": "ed25519", "public_key": base64.StdEncoding.EncodeToString(publicKey), "allowed_products": []string{testProduct}, "allowed_profiles": []string{testProfile}}
	writeJSON(t, filepath.Join(options.TrustDirectory, testKeyID+".json"), anchor)
	manager, err := NewManager(options)
	if err != nil {
		t.Fatal(err)
	}
	return manager, privateKey
}

func testManagerOptions(base string) ManagerOptions {
	return ManagerOptions{
		UpdateRoot: filepath.Join(base, "update"), PolicyPath: filepath.Join(base, "storage-layout.json"), TrustDirectory: filepath.Join(base, "trust"),
		Product: testProduct, Profile: testProfile, CurrentVersion: "1.0.0", Clock: func() time.Time { return time.Date(2026, 7, 23, 9, 0, 0, 0, time.UTC) },
	}
}

func signedPackage(t *testing.T, privateKey ed25519.PrivateKey, profile string, payload []byte) []byte {
	t.Helper()
	digest := sha256.Sum256(payload)
	manifest := Manifest{FormatVersion: 1, PackageID: "test-package-1", Product: testProduct, HardwareProfile: profile, Version: "1.1.0", MinimumCurrentVersion: "1.0.0", CreatedAt: time.Date(2026, 7, 23, 8, 0, 0, 0, time.UTC), SigningKeyID: testKeyID}
	manifest.Payloads = []Payload{{Name: "application", Path: "payload/application.bin", Class: "application", Size: int64(len(payload)), SHA256: hex.EncodeToString(digest[:])}}
	manifest.Apply.Capability = "validation-only"
	manifest.Recovery.Capability = "external-service"
	manifest.Recovery.Instructions = "Upload a known-good signed package; running firmware remains unchanged."
	manifestBytes, err := json.Marshal(manifest)
	if err != nil {
		t.Fatal(err)
	}
	envelope, err := json.Marshal(signatureEnvelope{Algorithm: "ed25519", KeyID: testKeyID, Signature: base64.StdEncoding.EncodeToString(ed25519.Sign(privateKey, manifestBytes))})
	if err != nil {
		t.Fatal(err)
	}
	var output bytes.Buffer
	writer := tar.NewWriter(&output)
	for _, entry := range []struct {
		name string
		data []byte
	}{{manifestPath, manifestBytes}, {signaturePath, envelope}, {"payload/application.bin", payload}} {
		if err := writer.WriteHeader(&tar.Header{Name: entry.name, Mode: 0o600, Size: int64(len(entry.data)), ModTime: time.Unix(0, 0)}); err != nil {
			t.Fatal(err)
		}
		if _, err := writer.Write(entry.data); err != nil {
			t.Fatal(err)
		}
	}
	if err := writer.Close(); err != nil {
		t.Fatal(err)
	}
	return output.Bytes()
}

func signedRecoveryPackage(t *testing.T, privateKey ed25519.PrivateKey, payload []byte) []byte {
	t.Helper()
	digest := sha256.Sum256(payload)
	manifest := Manifest{FormatVersion: 1, PackageID: "test-recovery-package-1", Product: testProduct, HardwareProfile: testProfile, Version: "1.1.0", MinimumCurrentVersion: "1.0.0", CreatedAt: time.Date(2026, 7, 23, 8, 0, 0, 0, time.UTC), SigningKeyID: testKeyID}
	manifest.Payloads = []Payload{{Name: "rootfs", Path: "payload/rootfs.img", Class: "full-rootfs", Size: int64(len(payload)), SHA256: hex.EncodeToString(digest[:]), TargetPartition: "rootfs", Encoding: "raw", ExpandedSize: int64(len(payload)), ExpandedSHA256: hex.EncodeToString(digest[:])}}
	manifest.Apply.Capability = ApplyCapabilityRecoveryV1
	manifest.Apply.RebootRequired = true
	manifest.Recovery.Capability = RecoveryCapabilityOfflineV1
	manifest.Recovery.Instructions = "Leave the request set and use COM5 recovery on failure."
	manifestBytes, err := json.Marshal(manifest)
	if err != nil {
		t.Fatal(err)
	}
	envelope, err := json.Marshal(signatureEnvelope{Algorithm: "ed25519", KeyID: testKeyID, Signature: base64.StdEncoding.EncodeToString(ed25519.Sign(privateKey, manifestBytes))})
	if err != nil {
		t.Fatal(err)
	}
	var output bytes.Buffer
	writer := tar.NewWriter(&output)
	for _, entry := range []struct {
		name string
		data []byte
	}{{manifestPath, manifestBytes}, {signaturePath, envelope}, {"payload/rootfs.img", payload}} {
		if err := writer.WriteHeader(&tar.Header{Name: entry.name, Mode: 0o600, Size: int64(len(entry.data)), ModTime: time.Unix(0, 0)}); err != nil {
			t.Fatal(err)
		}
		if _, err := writer.Write(entry.data); err != nil {
			t.Fatal(err)
		}
	}
	if err := writer.Close(); err != nil {
		t.Fatal(err)
	}
	return output.Bytes()
}

func writeJSON(t *testing.T, path string, value any) {
	t.Helper()
	data, err := json.Marshal(value)
	if err != nil {
		t.Fatal(err)
	}
	if err := os.WriteFile(path, data, 0o600); err != nil {
		t.Fatal(err)
	}
}

func TestFactoryResetErasesUploadedFirmwareAndKeepsInstalledVersion(t *testing.T) {
	m, key := newTestManager(t)
	payload := signedPackage(t, key, testProfile, []byte("customer-upload"))
	if _, err := m.Upload(context.Background(), UploadRequest{Filename: "release.bin", ContentLength: int64(len(payload)), Body: bytes.NewReader(payload)}); err != nil {
		t.Fatal(err)
	}
	if err := m.FactoryReset(context.Background()); err != nil {
		t.Fatal(err)
	}
	for _, dir := range []string{m.incoming, m.validated} {
		entries, err := os.ReadDir(dir)
		if err != nil || len(entries) != 0 {
			t.Fatal("uploaded firmware survived reset")
		}
	}
	reopened, err := NewManager(testManagerOptions(filepath.Dir(m.root)))
	if err != nil {
		t.Fatal(err)
	}
	status, err := reopened.Status(context.Background())
	if err != nil || status.State != StateIdle || status.DeviceVersion != m.currentVersion || status.Manifest != nil {
		t.Fatal("firmware reset state is incorrect")
	}
	m.active = true
	if err := m.FactoryReset(context.Background()); err == nil {
		t.Fatal("reset accepted during update")
	}
}

func TestManagerRecoversFieldApplyPolicyJournal(t *testing.T) {
	manager, _ := newTestManager(t)
	data, err := os.ReadFile(manager.journalPath)
	if err != nil {
		t.Fatal(err)
	}
	var value map[string]any
	if err := json.Unmarshal(data, &value); err != nil {
		t.Fatal(err)
	}
	value["apply_policy"] = "field"
	data, err = json.Marshal(value)
	if err != nil {
		t.Fatal(err)
	}
	if err := os.WriteFile(manager.journalPath, data, 0600); err != nil {
		t.Fatal(err)
	}
	reopened, err := NewManager(testManagerOptions(filepath.Dir(manager.root)))
	if err != nil {
		t.Fatal(err)
	}
	status, err := reopened.Status(context.Background())
	if err != nil || status.ApplyPolicy != "field" {
		t.Fatal("field policy was not preserved")
	}
}