package hostclaim

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"testing"
	"time"

	"github.com/carlosframework/platform/internal/release"
)

func retryAttempt(id string, session uint64) WakeRetryAttempt {
	return WakeRetryAttempt{ID: id, Box: "adi-knitting", Enrollment: "enroll-1", BootID: "boot-1",
		ControllerID: "controller-1", Session: session, InvocationID: "invoke-" + id, LaunchInputs: "sha256-inputs-" + id}
}

func TestWakeRetryDurableBudgetAndBackoff(t *testing.T) {
	ctx := context.Background()
	s := release.NewDirStore(t.TempDir())
	o := servingClaim(t, ctx, s)
	inc := o.Claim.Holder.Incarnation
	now := time.Date(2026, 10, 6, 12, 0, 0, 0, time.UTC)

	for n := 1; n <= WakeRetryLimit; n++ {
		a := retryAttempt(string(rune('0'+n)), o.Claim.Holder.Session)
		var err error
		o, err = Update(ctx, s, o, ReserveWakeRetryAttempt(WakeRetryReservation{
			HolderIncarnation: inc, EpisodeID: "episode-1", Attempt: a, Now: now,
		}))
		if err != nil {
			t.Fatalf("reserve %d: %v", n, err)
		}
		if o.Claim.Holder.WakeRetry.AttemptsReserved != n || o.Claim.Holder.WakeRetry.State != "attempting" {
			t.Fatalf("reserved %d = %+v", n, o.Claim.Holder.WakeRetry)
		}
		// A reread proves the budget is stored in the claim rather than process memory.
		if o, err = Read(ctx, s, testHost); err != nil {
			t.Fatal(err)
		}
		o, err = Update(ctx, s, o, FailWakeRetryAttempt(WakeRetryResult{
			HolderIncarnation: inc, EpisodeID: "episode-1", Attempt: a, Now: now,
		}))
		if err != nil {
			t.Fatalf("fail %d: %v", n, err)
		}
		wantState := "open"
		if n == WakeRetryLimit {
			wantState = "failed"
		}
		if r := o.Claim.Holder.WakeRetry; r.State != wantState || !r.NotBefore.Equal(now.Add(wakeRetryDelays[n-1])) {
			t.Fatalf("settled %d = %+v", n, r)
		}
		if _, err := Update(ctx, s, o, ReserveWakeRetryAttempt(WakeRetryReservation{HolderIncarnation: inc, Attempt: retryAttempt("early", 0), Now: now})); !errors.Is(err, ErrWrongState) {
			t.Fatalf("early reserve %d = %v", n, err)
		}
		now = now.Add(wakeRetryDelays[n-1])
	}
	if _, err := Update(ctx, s, o, ReserveWakeRetryAttempt(WakeRetryReservation{HolderIncarnation: inc, Attempt: retryAttempt("six", 0), Now: now})); !errors.Is(err, ErrWrongState) {
		t.Fatalf("sixth reserve = %v", err)
	}
	if _, err := Update(ctx, s, o, AdmitWriter(Admission{Region: "ie", PredecessorRegion: "ie", Box: "adi-knitting", Mode: "cold", NewGeneration: "new-after-failure", PredecessorCertified: true})); !errors.Is(err, ErrWrongState) {
		t.Fatalf("serving admission after exhausted budget = %v", err)
	}
}

func TestWakeRetryColdAdmissionBindingAndExactHealth(t *testing.T) {
	ctx := context.Background()
	s := release.NewDirStore(t.TempDir())
	o := servingClaim(t, ctx, s)
	inc := o.Claim.Holder.Incarnation
	a := retryAttempt("first", 0)
	now := time.Date(2026, 10, 6, 13, 0, 0, 0, time.UTC)
	var err error
	o, err = Update(ctx, s, o, ReserveWakeRetryAttempt(WakeRetryReservation{HolderIncarnation: inc, EpisodeID: "episode", Attempt: a, Now: now}))
	if err != nil {
		t.Fatal(err)
	}
	o, err = Update(ctx, s, o, AdmitWriter(Admission{Region: "ie", PredecessorRegion: "ie", Box: "adi-knitting", Mode: "cold", NewGeneration: "generation-1", PredecessorCertified: true}))
	if err != nil {
		t.Fatal(err)
	}
	session := o.Claim.Holder.Session
	if o.Claim.Holder.WakeRetry.AttemptsReserved != 1 { // cold admission did not reset the episode
		t.Fatalf("episode after admission = %+v", o.Claim.Holder.WakeRetry)
	}
	if _, err := Update(ctx, s, o, ClearWakeRetryHealthy(WakeRetryHealth{WakeRetryResult: WakeRetryResult{HolderIncarnation: inc, EpisodeID: "episode", Attempt: a}, ReadinessPassed: true})); !errors.Is(err, ErrWrongState) {
		t.Fatalf("health before session binding = %v", err)
	}
	o, err = Update(ctx, s, o, BindWakeRetrySession(inc, "episode", "first", session))
	if err != nil {
		t.Fatal(err)
	}
	a.Session = session
	wrong := a
	wrong.InvocationID = "other-invocation"
	if _, err := Update(ctx, s, o, ClearWakeRetryHealthy(WakeRetryHealth{WakeRetryResult: WakeRetryResult{HolderIncarnation: inc, EpisodeID: "episode", Attempt: wrong}, ReadinessPassed: true})); !errors.Is(err, ErrWrongState) {
		t.Fatalf("wrong invocation health = %v", err)
	}
	if _, err := Update(ctx, s, o, ClearWakeRetryHealthy(WakeRetryHealth{WakeRetryResult: WakeRetryResult{HolderIncarnation: inc, EpisodeID: "episode", Attempt: a}})); !errors.Is(err, ErrWrongState) {
		t.Fatalf("failed readiness = %v", err)
	}
	o, err = Update(ctx, s, o, ClearWakeRetryHealthy(WakeRetryHealth{WakeRetryResult: WakeRetryResult{HolderIncarnation: inc, EpisodeID: "episode", Attempt: a}, ReadinessPassed: true}))
	if err != nil || o.Claim.Holder.WakeRetry != nil {
		t.Fatalf("healthy clear = %+v, %v", o.Claim.Holder.WakeRetry, err)
	}
}

func TestWakeRetryOperatorResetRequiresSettledAttempt(t *testing.T) {
	ctx := context.Background()
	s := release.NewDirStore(t.TempDir())
	o := servingClaim(t, ctx, s)
	inc := o.Claim.Holder.Incarnation
	a := retryAttempt("active", o.Claim.Holder.Session)
	var err error
	o, err = Update(ctx, s, o, ReserveWakeRetryAttempt(WakeRetryReservation{HolderIncarnation: inc, EpisodeID: "old", Attempt: a, Now: time.Now()}))
	if err != nil {
		t.Fatal(err)
	}
	reset := WakeRetryReset{HolderIncarnation: inc, ObservedEpisode: "old", NewEpisodeID: "new", OperatorRequest: "request-42"}
	if _, err := Update(ctx, s, o, ResetWakeRetry(reset)); !errors.Is(err, ErrWrongState) {
		t.Fatalf("unsettled reset = %v", err)
	}
	reset.AttemptSettled = true
	o, err = Update(ctx, s, o, ResetWakeRetry(reset))
	if err != nil {
		t.Fatal(err)
	}
	if got := o.Claim.Holder.WakeRetry; got.EpisodeID != "new" || got.State != "open" || got.AttemptsReserved != 0 || got.ResetRequest != "request-42" {
		t.Fatalf("reset episode = %+v", got)
	}
}

func TestWakeRetryValidationAndWireShape(t *testing.T) {
	ctx := context.Background()
	s := release.NewDirStore(t.TempDir())
	o := servingClaim(t, ctx, s)
	bad := []WakeRetry{
		{EpisodeID: "e", State: "open", AttemptsReserved: 5},
		{EpisodeID: "e", State: "attempting", AttemptsReserved: 1},
		{EpisodeID: "e", State: "failed", AttemptsReserved: 4, NotBefore: time.Now()},
		{EpisodeID: "e", State: "unknown"},
	}
	for i := range bad {
		if _, err := Update(ctx, s, o, func(c *Claim) error { r := bad[i]; c.Holder.WakeRetry = &r; return nil }); !errors.Is(err, ErrMalformed) {
			t.Errorf("bad %d = %v", i, err)
		}
	}
	raw, err := json.Marshal(o.Claim.Holder)
	if err != nil {
		t.Fatal(err)
	}
	if bytes.Contains(raw, []byte(`"wake_retry"`)) {
		t.Fatalf("legacy holder gained wake_retry: %s", raw)
	}
}