package retention

import (
	"bufio"
	"context"
	"encoding/json"
	"errors"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"testing"
	"time"

	core "git.yuntuo.hk/k2safe/orbit/model"
)

func testJobSpec() JobSpec {
	return JobSpec{NamespaceID: "11111111111111111111111111111111", Kind: "tasks", PolicyRevision: 1, BatchSize: 1}
}
func testJobOptions() JobOptions { return JobOptions{MaxBatches: 10, MaxBytes: 16 << 20} }
func testJobSnapshot(t *testing.T, cursor string, payload int) ArchiveSnapshot {
	t.Helper()
	p := Defaults(testJobSpec().NamespaceID)
	p.Enabled = true
	p.Revision = 1
	next := "next-page"
	if cursor != "" {
		next = ""
	}
	raw, err := json.Marshal(map[string]any{"policy": p, "kind": "tasks", "cutoff": time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC), "root_count": 1, "source_retained": true, "next_cursor": next, "payload": strings.Repeat("x", payload)})
	if err != nil {
		t.Fatal(err)
	}
	return ArchiveSnapshot{1, digest(raw), string(raw)}
}
func testJobExporter(t *testing.T, calls *int) ExportBatch {
	return func(ctx context.Context, job core.ID, in ArchiveInput) (ArchiveSnapshot, error) {
		*calls++
		if job.Validate() != nil || in.ExpectedRevision != 1 || in.Kind != "tasks" || in.Limit != 1 {
			t.Fatal("export request changed identity")
		}
		return testJobSnapshot(t, in.Cursor, 0), nil
	}
}

func TestArchiveJobResumeAndOfflineVerification(t *testing.T) {
	dir := filepath.Join(t.TempDir(), "job")
	calls := 0
	options := testJobOptions()
	options.MaxBatches = 1
	first, err := RunArchive(context.Background(), dir, testJobSpec(), options, testJobExporter(t, &calls))
	if err != nil || first.Completed || first.Roots != 1 || first.Batches != 1 {
		t.Fatalf("first=%+v err=%v", first, err)
	}
	// 模拟强杀前只写了一部分临时文件；它不是有效检查点。
	if err := os.WriteFile(filepath.Join(dir, ".write.tmp"), []byte(`{"unfinished":`), 0600); err != nil {
		t.Fatal(err)
	}
	second, err := RunArchive(context.Background(), dir, testJobSpec(), options, testJobExporter(t, &calls))
	if err != nil || !second.Completed || second.Roots != 2 || second.Batches != 2 || second.JobID != first.JobID || calls != 2 {
		t.Fatalf("resume=%+v calls=%d err=%v", second, calls, err)
	}
	verified, err := VerifyArchive(context.Background(), dir)
	if err != nil || verified != second {
		t.Fatalf("verify=%+v err=%v", verified, err)
	}
	repeated, err := RunArchive(context.Background(), dir, testJobSpec(), options, testJobExporter(t, &calls))
	if err != nil || repeated != second || calls != 2 {
		t.Fatalf("completed job re-exported: %+v err=%v", repeated, err)
	}
	spec := testJobSpec()
	spec.PolicyRevision++
	if _, err := RunArchive(context.Background(), dir, spec, options, testJobExporter(t, &calls)); err == nil || calls != 2 {
		t.Fatal("changed scope accepted")
	}
	for _, name := range []string{"manifest.json", "batch-000001.json", "batch-000002.json"} {
		info, err := os.Stat(filepath.Join(dir, name))
		if err != nil || info.Mode().Perm() != 0600 {
			t.Fatalf("file permissions %s: %v", name, err)
		}
	}
}

func TestArchiveJobRejectsDamagedChain(t *testing.T) {
	for _, kind := range []string{"body", "missing", "swapped", "symlink", "unexpected"} {
		t.Run(kind, func(t *testing.T) {
			dir := filepath.Join(t.TempDir(), "job")
			calls := 0
			if _, err := RunArchive(context.Background(), dir, testJobSpec(), testJobOptions(), testJobExporter(t, &calls)); err != nil {
				t.Fatal(err)
			}
			first := filepath.Join(dir, "batch-000001.json")
			second := filepath.Join(dir, "batch-000002.json")
			var err error
			switch kind {
			case "body":
				raw, _ := os.ReadFile(first)
				var packet JobBatch
				if err = json.Unmarshal(raw, &packet); err != nil {
					t.Fatal(err)
				}
				packet.Snapshot.Content += " "
				raw, _ = json.Marshal(packet)
				err = os.WriteFile(first, raw, 0600)
			case "missing":
				err = os.Remove(first)
			case "swapped":
				raw, _ := os.ReadFile(second)
				err = os.WriteFile(first, raw, 0600)
			case "symlink":
				if err = os.Remove(first); err == nil {
					err = os.Symlink(second, first)
				}
			case "unexpected":
				err = os.WriteFile(filepath.Join(dir, "unrelated.json"), []byte("{}"), 0600)
			}
			if err != nil {
				t.Fatal(err)
			}
			if _, err := VerifyArchive(context.Background(), dir); err == nil {
				t.Fatal("damaged archive accepted")
			}
			if _, err := RunArchive(context.Background(), dir, testJobSpec(), testJobOptions(), testJobExporter(t, &calls)); err == nil || calls != 2 {
				t.Fatalf("resume skipped corrupt history: calls=%d err=%v", calls, err)
			}
		})
	}
}

func TestArchiveJobQuotaFailureAndCancellationPreserveProgress(t *testing.T) {
	dir := filepath.Join(t.TempDir(), "quota")
	options := testJobOptions()
	options.MaxBytes = 1 << 20
	calls := 0
	export := func(ctx context.Context, id core.ID, in ArchiveInput) (ArchiveSnapshot, error) {
		calls++
		return testJobSnapshot(t, in.Cursor, 700<<10), nil
	}
	status, err := RunArchive(context.Background(), dir, testJobSpec(), options, export)
	if err == nil || status.Completed || status.Batches != 1 {
		t.Fatalf("quota=%+v err=%v", status, err)
	}
	first, _ := os.ReadFile(filepath.Join(dir, "batch-000001.json"))
	options.MaxBytes = 4 << 20
	final, err := RunArchive(context.Background(), dir, testJobSpec(), options, export)
	if err != nil || !final.Completed || final.Batches != 2 || calls != 3 {
		t.Fatalf("quota resume=%+v calls=%d err=%v", final, calls, err)
	}
	after, _ := os.ReadFile(filepath.Join(dir, "batch-000001.json"))
	if string(first) != string(after) {
		t.Fatal("prior file overwritten")
	}
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	status, err = RunArchive(ctx, filepath.Join(t.TempDir(), "cancel"), testJobSpec(), testJobOptions(), func(ctx context.Context, id core.ID, in ArchiveInput) (ArchiveSnapshot, error) {
		cancel()
		return testJobSnapshot(t, "", 0), nil
	})
	if !errors.Is(err, context.Canceled) || status.Batches != 0 {
		t.Fatalf("cancellation advanced progress: %+v %v", status, err)
	}
}

func TestArchiveJobExclusiveLock(t *testing.T) {
	dir := filepath.Join(t.TempDir(), "job")
	root, lock, err := openJob(dir, true)
	if err != nil {
		t.Fatal(err)
	}
	defer root.Close()
	defer lock.Close()
	calls := 0
	if _, err := RunArchive(context.Background(), dir, testJobSpec(), testJobOptions(), testJobExporter(t, &calls)); err == nil || calls != 0 {
		t.Fatal("parallel writer accepted")
	}
}

func TestArchiveJobProcessKillAndResume(t *testing.T) {
	if dir := os.Getenv("ORBIT_ARCHIVE_KILL_TEST_DIR"); dir != "" {
		_, err := RunArchive(context.Background(), dir, testJobSpec(), testJobOptions(), func(ctx context.Context, id core.ID, in ArchiveInput) (ArchiveSnapshot, error) {
			if in.Cursor != "" {
				os.Stdout.WriteString("first-batch-durable\n")
				select {}
			}
			return testJobSnapshot(t, "", 0), nil
		})
		if err != nil {
			t.Fatal(err)
		}
		return
	}
	dir := filepath.Join(t.TempDir(), "job")
	ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
	defer cancel()
	cmd := exec.CommandContext(ctx, os.Args[0], "-test.run=^TestArchiveJobProcessKillAndResume$")
	cmd.Env = append(os.Environ(), "ORBIT_ARCHIVE_KILL_TEST_DIR="+dir)
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		t.Fatal(err)
	}
	cmd.Stderr = os.Stderr
	if err := cmd.Start(); err != nil {
		t.Fatal(err)
	}
	defer func() {
		if cmd.ProcessState == nil {
			cmd.Process.Kill()
			cmd.Wait()
		}
	}()
	scanner := bufio.NewScanner(stdout)
	if !scanner.Scan() || scanner.Text() != "first-batch-durable" {
		t.Fatal("child did not publish first batch")
	}
	if err := cmd.Process.Kill(); err != nil {
		t.Fatal(err)
	}
	if err := cmd.Wait(); err == nil {
		t.Fatal("child was not killed")
	}
	// 等待旧进程真正退出后，新进程重新取得锁，只导出第二批。
	calls := 0
	status, err := RunArchive(context.Background(), dir, testJobSpec(), testJobOptions(), testJobExporter(t, &calls))
	if err != nil || !status.Completed || status.Roots != 2 || status.Batches != 2 || calls != 1 {
		t.Fatalf("killed-process recovery=%+v calls=%d err=%v", status, calls, err)
	}
}

func TestArchiveJobWriteFailureDoesNotAdvance(t *testing.T) {
	if os.Geteuid() == 0 {
		t.Skip("permission failure requires an unprivileged process")
	}
	dir := filepath.Join(t.TempDir(), "job")
	calls := 0
	options := testJobOptions()
	options.MaxBatches = 1
	first, err := RunArchive(context.Background(), dir, testJobSpec(), options, testJobExporter(t, &calls))
	if err != nil {
		t.Fatal(err)
	}
	if err := os.Chmod(dir, 0500); err != nil {
		t.Fatal(err)
	}
	defer os.Chmod(dir, 0700)
	after, err := RunArchive(context.Background(), dir, testJobSpec(), options, testJobExporter(t, &calls))
	if err == nil || after != first {
		t.Fatalf("failed write advanced checkpoint: %+v err=%v", after, err)
	}
	verified, err := VerifyArchive(context.Background(), dir)
	if err != nil || verified != first {
		t.Fatalf("saved archive changed: %+v err=%v", verified, err)
	}
	if err := os.Chmod(dir, 0700); err != nil {
		t.Fatal(err)
	}
	final, err := RunArchive(context.Background(), dir, testJobSpec(), options, testJobExporter(t, &calls))
	if err != nil || !final.Completed || final.Roots != 2 {
		t.Fatalf("write recovery=%+v err=%v", final, err)
	}
}