//go:build linux

package host

import (
	"bytes"
	"context"
	"database/sql"
	"encoding/binary"
	"encoding/json"
	"errors"
	"fmt"
	"github.com/carlosframework/platform/internal/activator"
	"github.com/carlosframework/platform/internal/cgroup"
	"github.com/carlosframework/platform/internal/genselect"
	"github.com/carlosframework/platform/internal/hostbox"
	"github.com/carlosframework/platform/internal/hostclaim"
	"github.com/carlosframework/platform/internal/hostguard"
	"github.com/carlosframework/platform/internal/lease/leasetest"
	"github.com/carlosframework/platform/internal/moves"
	"github.com/carlosframework/platform/internal/park"
	"github.com/carlosframework/platform/internal/provision"
	"github.com/carlosframework/platform/internal/registry"
	"github.com/carlosframework/platform/internal/release"
	"github.com/superfly/ltx"
	"golang.org/x/sys/unix"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"testing"
)

func TestPlan12RecoveryRestoreSourceDatabase(t *testing.T) {
	t.Parallel()
	for _, name := range []string{hostclaim.KindRestoring, hostclaim.KindReverting, "clean-return", "never-restoring", "never-reverting", "never-clean-return", "missing-inventory", "wrong-inventory", "missing-registry", "interrupted", "admission-crash", "prejournal-failure", "intent-crash", "evidence-read-failure", "concurrent-update", "prelaunch-update"} {
		t.Run(name, func(t *testing.T) {
			kind, clean := name, name == "clean-return" || name == "never-clean-return"
			bad := name == "missing-inventory" || name == "wrong-inventory" || name == "missing-registry"
			beforeAdmission := name == "prejournal-failure" || name == "intent-crash" || name == "evidence-read-failure"
			retryCase := name == "interrupted" || name == "admission-crash" || beforeAdmission || name == "prelaunch-update"
			if bad || retryCase || name == "concurrent-update" {
				kind = hostclaim.KindRestoring
			}
			if name == "never-restoring" {
				kind = hostclaim.KindRestoring
			}
			if name == "never-reverting" {
				kind = hostclaim.KindReverting
			}
			never := name == "never-restoring" || name == "never-reverting" || name == "never-clean-return"
			if clean {
				kind = hostclaim.KindRestoring
			}
			f, o, h := plan12Handover(t, hostclaim.KindRestoring, "stop")
			stamp := o.Claim.Holder.Stamp
			baseline := plan12RestoreBaseline(t, f)
			db := provision.NestedDBPath(f.guard.DataDir, o.Host)
			dbRel, _ := filepath.Rel(f.guard.DataDir, db)
			dirRel, _ := filepath.Rel(f.guard.DataDir, filepath.Dir(db))
			if err := os.MkdirAll(filepath.Dir(db), 0700); err != nil {
				t.Fatal(err)
			}
			if err := os.Mkdir(filepath.Join(f.guard.DataDir, ".parked"), 0770); err != nil {
				t.Fatal(err)
			}
			if err := os.WriteFile(db, []byte("source database"), 0600); err != nil {
				t.Fatal(err)
			}
			if err := hostbox.WriteGenerationMarker(db+".gen", "src"); err != nil {
				t.Fatal(err)
			}
			key := moves.StoppedKey(moves.ReportPrefix("bab", "jam", o.Host, "m1"), "source")
			report := plan12Stopped(o, "source", "box-1")
			raw, _ := json.Marshal(report)
			if _, err := f.claims.Put(t.Context(), key, raw, ""); err != nil {
				t.Fatal(err)
			}
			var err error
			o, err = hostclaim.Update(t.Context(), f.claims, o, func(c *hostclaim.Claim) error {
				c.Op.Kind = kind
				c.Op.Journal[hostclaim.JournalSourceStamp] = strconv.FormatUint(stamp, 10)
				c.Op.Journal[hostclaim.JournalSrcGen] = "src"
				c.Op.Journal[hostclaim.JournalCutSession] = "2"
				c.Op.Journal[hostclaim.JournalEvidenceSource] = key
				if kind == hostclaim.KindRestoring {
					c.Holder.Admissions = []hostclaim.AdmissionRecord{{Session: 2, Box: "box-2", Generation: "target-gen", Mode: "cold", Region: "uk", Predecessor: "src", PredecessorRegion: "ie"}}
					if clean {
						c.Holder.Generation = "src"
					}
				}
				if kind == hostclaim.KindReverting {
					c.Holder.WriterBox = "box-1"
					c.Holder.Session = 2
					c.Holder.Generation = "src"
					c.Holder.Admissions = []hostclaim.AdmissionRecord{{Session: 2, Box: "box-1", Generation: "src", Mode: "cold", Region: "ie"}}
				}
				if never {
					c.Holder.WriterBox = ""
					c.Holder.Session = 0
					c.Holder.Generation = ""
					c.Holder.Admissions = nil
					c.Op.Journal[hostclaim.JournalCutSession] = "0"
					c.Op.Journal[hostclaim.JournalSrcGen] = ""
				}
				return nil
			})
			if err != nil {
				t.Fatal(err)
			}
			if never {
				if err = os.Remove(db); err != nil {
					t.Fatal(err)
				}
				if err = os.Remove(db + ".gen"); err != nil {
					t.Fatal(err)
				}
			}
			local := park.Local{Guard: f.guard, Claims: f.claims, DataDir: f.guard.DataDir, Box: "box-1"}
			if err = local.RecordInventory(t.Context(), o.Host, hostguard.Captured{Incarnation: o.Claim.Holder.Incarnation, Stamp: stamp}, "flat", []string{dbRel, dbRel + ".gen"}, func(*hostclaim.Claim) error { return nil }); err != nil {
				t.Fatal(err)
			}
			o = f.read()
			c := f.controller("boot-1", "controller-1", 4101, "start-1")
			if claimInitialSourceRecovery(o, "box-1") {
				if never {
					f.kernel.noProcess = true
				} else {
					target := recoveryTuple(o, c.local)
					if err := f.journal.Prelaunch(t.Context(), hostbox.LaunchInventory{ControllerPID: 4101, ProcessStart: "start-1", Target: target, Groups: []hostbox.LaunchGroup{{RelativePath: "carlos/" + claimContainmentName(target), Device: 1, Inode: target.Session}}, NoState: true}); err != nil {
						t.Fatal(err)
					}
				}
			}
			c.wake.Inputs = plan12ReturnInputs{plan9D2Inputs: plan9D2Inputs{recoveryInputs{f}}, clean: clean, baseline: baseline}
			c.wake.Continue = hostclaim.SessionContinues
			c.wake.Operations = &plan9D2Attempts{}
			c.wake.ProbeReady = func(context.Context, *exec.Cmd, registry.Route, ClaimLaunchInputs, string) error { return nil }
			retirement := claimArtifactRetirement{Local: local, Continue: hostclaim.SessionContinues}
			c.wake.RetireArtifacts, c.wake.RecoverArtifacts = retirement.Retire, retirement.Recover
			c.assembly.Guard = f.guard
			c.assembly.bindClaimWakeStop()
			inv, e := h.Inventory.ReadInstallation(t.Context(), o.Host, o.Claim.Holder.Incarnation, stamp)
			if e != nil {
				t.Fatal(e)
			}
			var stat unix.Stat_t
			if e = unix.Lstat(filepath.Dir(db), &stat); e != nil {
				t.Fatal(e)
			}
			inv.Entries = append(inv.Entries, hostbox.StopInventoryEntry{Identity: inventoryIdentityFixture(t, f.guard.DataDir, dirRel), Kind: "path", ArtifactID: "state", ManagedRoot: "data", RelativePath: dirRel, Ownership: "old-owner", Device: uint64(stat.Dev), Inode: stat.Ino, Type: unix.S_IFDIR, Materialization: "materialized"})
			if e = h.Inventory.WriteStopInventory(t.Context(), inv); e != nil {
				t.Fatal(e)
			}
			roles := &ClaimRoleAssembly{Manager: c.wake.Manager, Journal: f.journal, Inventory: h.Inventory, ManagedRoots: map[string]string{"data": f.guard.DataDir}, Identity: h.Identity, BootID: "boot-1", ControllerID: "controller-1"}
			roles.prepareGroup = func(name string, _ cgroup.Limit) (*cgroup.Prepared, hostbox.LaunchGroup, error) {
				return nil, hostbox.LaunchGroup{RelativePath: "carlos/" + name, Device: 1, Inode: 1}, nil
			}
			c.wake.Prelaunch, c.wake.PrepareState = roles.Prelaunch, roles.PrepareState
			if name == "missing-inventory" {
				replacement, e := openLegacyInventory(t, filepath.Join(t.TempDir(), "empty.db"))
				if e != nil {
					t.Fatal(e)
				}
				defer replacement.Close()
				h.Inventory = replacement
			}
			if name == "wrong-inventory" {
				inv.Enrollment = "other"
				if e = h.Inventory.WriteStopInventory(t.Context(), inv); e != nil {
					t.Fatal(e)
				}
			}
			if name == "missing-registry" {
				h.Registry = nil
			}
			h.Launches = f.journal
			savedRegistry := h.Registry
			interrupted := false
			if name == "interrupted" {
				h.afterRecoveryRows = func() error { interrupted = true; return hostbox.ErrChanged }
				if _, e := h.Registry.UpsertTimerRegistration(registry.Timer{Host: o.Host, Name: "uncensused", Kind: "cron", Spec: "* * * * *", ClaimOwnership: registry.ClaimOwnership{ClaimIncarnation: o.Claim.Holder.Incarnation, ClaimStamp: stamp, ClaimOwner: "old-owner", ClaimRowID: "timer-row"}}, ""); e != nil {
					t.Fatal(e)
				}
			}
			wireClaimAdmission(c.wake, "box-1", c.wake.AfterAdmission, h, nil, nil)
			if name == "concurrent-update" {
				h.afterRecoveryRows = func() error {
					_, err := hostclaim.Update(t.Context(), f.claims, f.read(), func(c *hostclaim.Claim) error { c.Op.Journal["concurrent"] = "acknowledged"; return nil })
					return err
				}
				wireClaimAdmission(c.wake, "box-1", c.wake.AfterAdmission, h, nil, nil)
			}
			if name == "intent-crash" {
				prepare := c.wake.recoveryIntent
				c.wake.recoveryIntent = func(ctx context.Context, o hostclaim.Observed, in hostclaim.Admission) (hostclaim.Admission, error) {
					in, err := prepare(ctx, o, in)
					if err != nil {
						return in, err
					}
					interrupted = true
					return in, hostbox.ErrChanged
				}
			}
			if name == "prelaunch-update" {
				c.wake.recoveryPrelaunched = func(context.Context, hostclaim.Observed) error { interrupted = true; return hostbox.ErrChanged }
			}
			if name == "admission-crash" {
				c.wake.AfterAdmission = func(context.Context, hostclaim.Observed, hostclaim.Observed) error {
					interrupted = true
					return hostbox.ErrChanged
				}
			}
			if name == "prejournal-failure" {
				h.Registry = nil
				wireClaimAdmission(c.wake, "box-1", c.wake.AfterAdmission, h, nil, nil)
			}
			if err := c.wake.stopPredecessor(t.Context(), o); err != nil {
				t.Fatal(err)
			}
			if name == "evidence-read-failure" {
				h.Evidence = &plan12FailEvidence{Store: f.claims, fail: true}
				wireClaimAdmission(c.wake, "box-1", c.wake.AfterAdmission, h, nil, nil)
			}
			live, cancel := context.WithCancel(t.Context())
			defer cancel()
			done := make(chan struct{})
			go func() { <-live.Done(); close(done) }()
			c.wake.mu.Lock()
			c.wake.pending[o.Host] = &claimWakeAttempt{observed: o, held: recoveryHeld{}, ctx: live, cancel: cancel, done: done}
			c.wake.mu.Unlock()
			admitted, err := c.runtime.Admit(t.Context(), registry.Route{Host: o.Host}, o, activator.StartPolicy{Probe: true})
			if retryCase {
				if name == "prejournal-failure" {
					interrupted = err != nil
					h.Registry = savedRegistry
				}
				if name == "evidence-read-failure" {
					interrupted = err != nil
					h.Evidence = f.claims
				}
				if err == nil || !interrupted {
					t.Fatal("did not interrupt after registry commit", err)
				}
				first := f.read()
				first, e = hostclaim.Update(t.Context(), f.claims, first, func(c *hostclaim.Claim) error {
					c.Op.Journal["unrelated"] = "park acknowledged"
					return nil
				})
				if e != nil {
					t.Fatal(e)
				}
				row, e := h.Registry.Get(o.Host)
				if e != nil || name == "interrupted" && row.ClaimStamp != first.Claim.Holder.Stamp {
					t.Fatal("registry did not commit", e)
				}
				if name == "interrupted" {
					if _, e = h.Inventory.ReadInstallation(t.Context(), o.Host, first.Claim.Holder.Incarnation, first.Claim.Holder.Stamp); !errors.Is(e, sql.ErrNoRows) {
						t.Fatal("successor inventory already written", e)
					}
				}
				c = f.controller("boot-2", "controller-2", 4102, "start-2")
				c.wake.Inputs = plan12ReturnInputs{plan9D2Inputs: plan9D2Inputs{recoveryInputs{f}}, baseline: baseline}
				c.wake.Continue = hostclaim.SessionContinues
				c.wake.Operations = &plan9D2Attempts{}
				c.wake.ProbeReady = func(context.Context, *exec.Cmd, registry.Route, ClaimLaunchInputs, string) error { return nil }
				c.wake.RetireArtifacts, c.wake.RecoverArtifacts = retirement.Retire, retirement.Recover
				c.wake.Prelaunch, c.wake.PrepareState = roles.Prelaunch, roles.PrepareState
				c.assembly.Guard = f.guard
				c.assembly.bindClaimWakeStop()
				h.afterRecoveryRows = nil
				wireClaimAdmission(c.wake, "box-1", c.wake.AfterAdmission, h, nil, nil)
				if e := c.wake.stopPredecessor(t.Context(), first); e != nil {
					t.Fatal("retry chose predecessor stop before handover", e)
				}
				if name == "prelaunch-update" {
					bodies, err := h.Registry.PendingRecoveryHandovers(t.Context())
					if err != nil || len(bodies) != 0 {
						t.Fatalf("prelaunch journal stayed pending: %+v %v", bodies, err)
					}
					if len(f.requests) == 0 {
						t.Fatal("prelaunch retry did not reach root stop")
					}
					return
				}
				retryCtx, retryCancel := context.WithCancel(t.Context())
				defer retryCancel()
				retryDone := make(chan struct{})
				go func() { <-retryCtx.Done(); close(retryDone) }()
				c.wake.pending[o.Host] = &claimWakeAttempt{observed: first, held: recoveryHeld{}, ctx: retryCtx, cancel: retryCancel, done: retryDone}
				admitted, err = c.runtime.Admit(t.Context(), registry.Route{Host: o.Host}, first, activator.StartPolicy{Probe: true})
				if err != nil || !beforeAdmission && admitted.Claim.Holder.Session != first.Claim.Holder.Session {
					t.Fatal("retry did not resume admitted session", err)
				}
				if !beforeAdmission {
					if err = c.runtime.ResumedAdmission(t.Context(), first, admitted); err != nil {
						t.Fatal("activator could not verify resumed admission", err)
					}
					if err := c.runtime.ResumedAdmission(t.Context(), first, admitted); !errors.Is(err, hostbox.ErrChanged) {
						t.Fatal("resume verification was not single use", err)
					}
				}
			}
			if bad {
				if err == nil {
					t.Fatal("unproved inventory admitted")
				}
				return
			}
			if err != nil {
				t.Fatal(err)
			}
			launch, err := c.runtime.Prepare(t.Context(), registry.Route{Host: o.Host}, admitted, activator.StartPolicy{Probe: true})
			if err != nil {
				t.Fatal(err)
			}
			if never {
				if f.kernel.noProcessCalls == 0 {
					t.Fatal("R4 witness was not checked")
				}
				return
			}
			restoreErr := launch.Restore(t.Context())
			if name == "concurrent-update" {
				if !errors.Is(restoreErr, hostclaim.ErrConflict) {
					t.Fatal("stale admission accepted foreign journal write", restoreErr)
				}
				if f.read().Claim.Op.Journal["concurrent"] != "acknowledged" {
					t.Fatal("concurrent journal write lost")
				}
				if _, err := os.Stat(db); !errors.Is(err, os.ErrNotExist) {
					t.Fatal("foreign update reached database installation", err)
				}
				retired := filepath.Join(f.guard.DataDir, ".parked", admitted.Claim.Holder.Incarnation, o.Host, fmt.Sprintf("superseded-%d-%d", admitted.Claim.Holder.Session, stamp), "artifacts", dbRel)
				if got, err := os.ReadFile(retired); err != nil || string(got) != "source database" {
					t.Fatalf("foreign update lost retired predecessor: %q %v", got, err)
				}
				return
			}
			if restoreErr != nil {
				t.Fatal(restoreErr)
			}
			if got, err := os.ReadFile(db + ".gen"); err != nil || string(got) != admitted.Claim.Holder.Generation+"\n" {
				t.Fatal(string(got), err)
			}
			restored, err := sql.Open("sqlite", "file:"+db+"?mode=ro")
			if err != nil {
				t.Fatal(err)
			}
			defer restored.Close()
			var value string
			if err = restored.QueryRow("SELECT value FROM fixture").Scan(&value); err != nil || value != "selected recovery baseline" {
				t.Fatalf("restored database: %q %v", value, err)
			}
			retired := filepath.Join(f.guard.DataDir, ".parked", admitted.Claim.Holder.Incarnation, o.Host, fmt.Sprintf("superseded-%d-%d", admitted.Claim.Holder.Session, stamp), "artifacts", dbRel)
			if got, err := os.ReadFile(retired); err != nil || string(got) != "source database" {
				t.Fatalf("retired predecessor: %q %v", got, err)
			}
			if got, err := os.ReadFile(retired + ".gen"); err != nil || string(got) != "src\n" {
				t.Fatalf("retired marker: %q %v", got, err)
			}
			if tag := f.read().Claim.Inventory["box-1"].Paths[dbRel]; tag != admitted.Claim.Holder.Stamp {
				t.Fatal("old ownership", tag)
			}
		})
	}
}

type plan12ReturnInputs struct {
	plan9D2Inputs
	clean    bool
	baseline *genselect.Plan
}

func (in plan12ReturnInputs) Select(ctx context.Context, rt registry.Route, o hostclaim.Observed, p activator.StartPolicy) (ClaimLaunchInputs, error) {
	selected, err := in.plan9D2Inputs.Select(ctx, rt, o, p)
	selected.Admission.KeepStamp = false
	if in.baseline != nil {
		selected.Empty = false
		selected.Baseline = in.baseline
		selected.Admission.RemoteMaxTXID = in.baseline.MaxTXID
	}
	selected.DBPath = provision.NestedDBPath(in.f.guard.DataDir, o.Host)
	if in.clean {
		selected.Admission.Coverage = "clean-return"
	}
	return selected, err
}

func plan12RestoreBaseline(t *testing.T, f *recoveryFixture, extraBytes ...int) *genselect.Plan {
	t.Helper()
	path := filepath.Join(t.TempDir(), "selected.db")
	db, err := sql.Open("sqlite", path)
	if err != nil {
		t.Fatal(err)
	}
	if _, err := db.Exec("CREATE TABLE fixture(value TEXT); INSERT INTO fixture VALUES ('selected recovery baseline')"); err != nil {
		t.Fatal(err)
	}
	if len(extraBytes) > 0 && extraBytes[0] > 0 {
		if _, err := db.Exec("CREATE TABLE padding(value BLOB); INSERT INTO padding VALUES (zeroblob(?))", extraBytes[0]); err != nil {
			t.Fatal(err)
		}
	}
	if err := db.Close(); err != nil {
		t.Fatal(err)
	}
	raw, err := os.ReadFile(path)
	if err != nil {
		t.Fatal(err)
	}
	size := int(binary.BigEndian.Uint16(raw[16:18]))
	if size == 1 {
		size = 65536
	}
	spec := ltx.FileSpec{Header: ltx.Header{Version: ltx.Version, Flags: ltx.HeaderFlagNoChecksum, PageSize: uint32(size), Commit: uint32(len(raw) / size), MinTXID: 1, MaxTXID: 1}}
	for off := 0; off < len(raw); off += size {
		pg := uint32(off/size + 1)
		if pg != ltx.LockPgno(uint32(size)) {
			spec.Pages = append(spec.Pages, ltx.PageSpec{Header: ltx.PageHeader{Pgno: pg}, Data: raw[off : off+size]})
		}
	}
	var out bytes.Buffer
	if _, err := spec.WriteTo(&out); err != nil {
		t.Fatal(err)
	}
	key := "instances/" + f.hostName() + "/g-src/0009/" + ltx.FormatFilename(1, 1)
	bucket := leasetest.NewVersionedBucket()
	bucket.AddBucket("data")
	version := bucket.Put("data", key, out.Bytes())
	bucket.Put("data", key, []byte("newer unselected version"))
	f.data, f.bucket = bucket, "data"
	return &genselect.Plan{Generation: "src", MaxTXID: 1, Pin: &genselect.Pin{SnapshotCRC64: uint64(spec.Trailer.FileChecksum)}, Files: []genselect.File{{Key: key, VersionID: version, Level: 9, MinTXID: 1, MaxTXID: 1, Size: int64(out.Len())}}}
}

type plan12FailEvidence struct {
	release.Store
	fail bool
}

func (s *plan12FailEvidence) Get(ctx context.Context, key string) ([]byte, string, error) {
	if s.fail {
		s.fail = false
		return nil, "", hostbox.ErrChanged
	}
	return s.Store.Get(ctx, key)
}