package runtime

import (
	"archive/tar"
	"bytes"
	"compress/gzip"
	"context"
	"errors"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"testing"
	"time"

	"browser.local/run/protocol"
)

func TestRunSelfUpdateResumesPartialDownloadAndRejectsChecksum(t *testing.T) {
	assignment, input, payload := selfUpdateTestFixture(t)
	client := newFakeWorkerClient()
	client.updateInput = input
	client.updatePayload = payload
	cfg := workerTestConfig(t)
	transactionRoot := filepath.Join(cfg.WorkspaceRoot, "self-updates", safeWorkspaceName(assignment.JobID))
	if err := os.MkdirAll(transactionRoot, 0o700); err != nil {
		t.Fatalf("create transaction root: %v", err)
	}
	partial := len(payload) / 3
	if err := os.WriteFile(filepath.Join(transactionRoot, "update.archive"), payload[:partial], 0o600); err != nil {
		t.Fatalf("write partial update: %v", err)
	}
	worker, err := NewWorker(cfg, client, WithSelfUpdateActivator(&recordingSelfUpdateActivator{}))
	if err != nil {
		t.Fatalf("new worker: %v", err)
	}
	worker.state.SessionToken = "session-token"
	result := worker.executeRunSelfUpdate(context.Background(), assignment)
	if result.State != lifecycleResultStateSucceeded || result.ExecutionResult.Kind != "run.update.staged" || result.ActivationManifest == "" {
		t.Fatalf("expected staged self-update, got %+v", result)
	}
	if len(client.updateChunkOffsets) == 0 || client.updateChunkOffsets[0] != int64(partial) {
		t.Fatalf("expected resumable range from %d, got %+v", partial, client.updateChunkOffsets)
	}

	badClient := newFakeWorkerClient()
	badInput := input
	badInput.Checksum = "sha256:bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb"
	badClient.updateInput = badInput
	badClient.updatePayload = payload
	badWorker, err := NewWorker(workerTestConfig(t), badClient, WithSelfUpdateActivator(&recordingSelfUpdateActivator{}))
	if err != nil {
		t.Fatalf("new bad checksum worker: %v", err)
	}
	badWorker.state.SessionToken = "session-token"
	bad := badWorker.executeRunSelfUpdate(context.Background(), assignment)
	if bad.State != lifecycleResultStateFailed || bad.ErrorCode != "self_update_download_failed" {
		t.Fatalf("expected final checksum rejection, got %+v", bad)
	}
}

func TestSelfUpdateActivationWaitsForAcceptedTerminalResult(t *testing.T) {
	assignment, input, payload := selfUpdateTestFixture(t)
	client := newFakeWorkerClient()
	client.claimJob = assignment
	client.updateInput = input
	client.updatePayload = payload
	client.resultErr = errors.New("stale lease")
	activator := &recordingSelfUpdateActivator{}
	worker, err := NewWorker(workerTestConfig(t), client, WithSelfUpdateActivator(activator))
	if err != nil {
		t.Fatalf("new worker: %v", err)
	}
	if err := worker.Register(context.Background()); err != nil {
		t.Fatalf("register: %v", err)
	}
	handled, err := worker.ClaimAndRunOnce(context.Background())
	if !handled || err == nil {
		t.Fatalf("expected rejected result error, handled=%v err=%v", handled, err)
	}
	if activator.manifestPath != "" {
		t.Fatalf("stale terminal result must not activate update: %s", activator.manifestPath)
	}
	if worker.journal.ActiveCount() != 1 {
		t.Fatal("staged result must remain recoverable until Platform accepts it")
	}
}

func TestSelfUpdateReplacementPreservesRollbackAndRejectsTraversal(t *testing.T) {
	root := t.TempDir()
	current := filepath.Join(root, "run")
	staged := filepath.Join(root, "staged-run")
	backup := filepath.Join(root, "previous-run")
	if err := os.WriteFile(current, []byte("old-run"), 0o700); err != nil {
		t.Fatalf("write current: %v", err)
	}
	if err := os.WriteFile(staged, []byte("new-run"), 0o700); err != nil {
		t.Fatalf("write staged: %v", err)
	}
	checksum, _, err := checksumFile(staged)
	if err != nil {
		t.Fatalf("checksum staged: %v", err)
	}
	manifest := SelfUpdateManifest{StagedExecutable: staged, CurrentExecutable: current, BackupExecutable: backup, BinaryChecksum: checksum}
	if err := replaceRunExecutable(manifest); err != nil {
		t.Fatalf("replace executable: %v", err)
	}
	if body, _ := os.ReadFile(current); string(body) != "new-run" {
		t.Fatalf("expected new executable, got %q", body)
	}
	if err := rollbackRunExecutable(manifest); err != nil {
		t.Fatalf("rollback executable: %v", err)
	}
	if body, _ := os.ReadFile(current); string(body) != "old-run" {
		t.Fatalf("expected previous executable after rollback, got %q", body)
	}

	archivePath := filepath.Join(root, "unsafe.tar.gz")
	var archive bytes.Buffer
	gzipWriter := gzip.NewWriter(&archive)
	tarWriter := tar.NewWriter(gzipWriter)
	body := []byte("escape")
	if err := tarWriter.WriteHeader(&tar.Header{Name: "../run", Mode: 0o700, Size: int64(len(body)), Typeflag: tar.TypeReg}); err != nil {
		t.Fatalf("write unsafe header: %v", err)
	}
	_, _ = tarWriter.Write(body)
	_ = tarWriter.Close()
	_ = gzipWriter.Close()
	if err := os.WriteFile(archivePath, archive.Bytes(), 0o600); err != nil {
		t.Fatalf("write unsafe archive: %v", err)
	}
	if _, err := extractRunUpdateBinary(archivePath, "tar.gz", "run", filepath.Join(root, "escaped")); err == nil || !strings.Contains(err.Error(), "unsafe") {
		t.Fatalf("expected traversal rejection, got %v", err)
	}
}

func TestPrepareSelfUpdateManifestRevalidatesStagedBinaryAfterRestart(t *testing.T) {
	root := t.TempDir()
	assignment := workerJobAssignment(protocol.RunCapabilityRunSelfUpdate)
	assignment.TargetKey = "run/update"
	assignment.InputRef = "artifact://artifact-run-staged"
	input := protocol.RunUpdateInputResponse{JobID: assignment.JobID, ServerInstanceID: assignment.ServerInstanceID, RunEndpointID: assignment.RunEndpointID, ArtifactID: "artifact-run-staged", Checksum: "sha256:" + strings.Repeat("a", 64), SizeBytes: 16, TargetOS: runtime.GOOS, TargetArch: runtime.GOARCH, PackageFormat: "tar.gz", ExecutableName: "run", TargetRelease: "release-staged", ChunkSizeBytes: 8}
	manifestPath := filepath.Join(root, "manifest.json")
	stagedPath := filepath.Join(root, "run.staged")
	if err := os.WriteFile(stagedPath, []byte("staged-binary"), 0o700); err != nil {
		t.Fatalf("write staged binary: %v", err)
	}
	checksum, _, err := checksumFile(stagedPath)
	if err != nil {
		t.Fatalf("checksum staged binary: %v", err)
	}
	manifest := SelfUpdateManifest{Version: selfUpdateManifestVersion, JobID: assignment.JobID, Attempt: assignment.Attempt, LeaseToken: assignment.LeaseToken, ArtifactID: input.ArtifactID, ArtifactChecksum: input.Checksum, ArtifactSizeBytes: input.SizeBytes, TargetOS: input.TargetOS, TargetArch: input.TargetArch, TargetRelease: input.TargetRelease, StagedExecutable: stagedPath, BinaryChecksum: checksum, Phase: "staged", CreatedAt: time.Now().UTC(), UpdatedAt: time.Now().UTC()}
	if err := persistSelfUpdateManifest(manifestPath, manifest); err != nil {
		t.Fatalf("persist staged manifest: %v", err)
	}
	if _, err := prepareSelfUpdateManifest(manifestPath, assignment, input, root); err != nil {
		t.Fatalf("revalidate staged manifest: %v", err)
	}
	if err := os.WriteFile(stagedPath, []byte("tampered-binary"), 0o700); err != nil {
		t.Fatalf("tamper staged binary: %v", err)
	}
	if _, err := prepareSelfUpdateManifest(manifestPath, assignment, input, root); err == nil || !strings.Contains(err.Error(), "checksum changed") {
		t.Fatalf("expected staged checksum rejection, got %v", err)
	}
	manifest.StagedExecutable = filepath.Join(root, "..", "outside")
	if err := persistSelfUpdateManifest(manifestPath, manifest); err != nil {
		t.Fatalf("persist unsafe staged manifest: %v", err)
	}
	if _, err := prepareSelfUpdateManifest(manifestPath, assignment, input, root); err == nil || !strings.Contains(err.Error(), "outside") {
		t.Fatalf("expected staged path rejection, got %v", err)
	}
}

func selfUpdateTestFixture(t *testing.T) (protocol.RunJobAssignment, protocol.RunUpdateInputResponse, []byte) {
	t.Helper()
	assignment := workerJobAssignment(protocol.RunCapabilityRunSelfUpdate)
	assignment.TargetKey = "run/update"
	assignment.InputRef = "artifact://artifact-run-latest"
	assignment.State = "running"
	assignment.LeaseToken = "lease-update"
	assignment.Attempt = 1
	executableName := "run"
	if runtime.GOOS == "windows" {
		executableName = "run.exe"
	}
	payload := []byte("self-update-test-binary")
	input := protocol.RunUpdateInputResponse{JobID: assignment.JobID, ServerInstanceID: assignment.ServerInstanceID, RunEndpointID: assignment.RunEndpointID, ArtifactID: "artifact-run-latest", Checksum: bytesChecksum(payload), SizeBytes: int64(len(payload)), TargetOS: runtime.GOOS, TargetArch: runtime.GOARCH, PackageFormat: "raw-executable", ExecutableName: executableName, TargetRelease: "run-release-test", ChunkSizeBytes: 64}
	return assignment, input, payload
}