package main

import (
	"encoding/json"
	"testing"
	"time"
)

func TestFailureKinds(t *testing.T) {
	for _, raw := range []string{
		`{"codexErrorInfo":"cyberPolicy"}`,
		`{"codexErrorInfo":"bioPolicy"}`,
		`{"codexErrorInfo":"misalignmentPolicyViolation"}`,
		`{"codexErrorInfo":"unauthorized"}`,
		`{"codexErrorInfo":"usageLimitExceeded"}`,
		`{"codexErrorInfo":"badRequest"}`,
		`{"codexErrorInfo":"other"}`,
		`{"codexErrorInfo":{"httpConnectionFailed":{"httpStatusCode":401}}}`,
		`{"codexErrorInfo":{"futureError":{"httpStatusCode":503}}}`,
		`{"codexErrorInfo":null}`,
		`{}`,
		`null`,
	} {
		t.Run(raw, func(t *testing.T) {
			var failure *turnError
			if err := json.Unmarshal([]byte(raw), &failure); err != nil {
				t.Fatal(err)
			}
			f := newFake(t, nil)
			if err := f.reconcile(); err != nil {
				t.Fatal(err)
			}
			w := f.op.store.watch("thread-1")
			for attempt := 1; attempt <= 2; attempt++ {
				f.fail()
				if attempt == 2 {
					f.turn.ID = "turn-2"
				}
				f.turn.Error = failure
				if err := f.reconcile(); err != nil {
					t.Fatal(err)
				}
				if !w.eligible() || w.Failures != attempt || f.count("thread/goal/set") != attempt-1 {
					t.Fatalf("failed turn did not enter backoff: %+v", w)
				}
				if delay := time.Until(w.Due); delay <= 0 || delay > time.Duration(attempt)*time.Second {
					t.Fatalf("unexpected backoff: %s", delay)
				}
				if err := f.reconcile(); err != nil {
					t.Fatal(err)
				}
				if f.count("thread/goal/set") != attempt-1 {
					t.Fatal("resumed before backoff elapsed")
				}
				w.Due = time.Now().Add(-time.Second)
				if err := f.reconcile(); err != nil {
					t.Fatal(err)
				}
				if f.count("thread/goal/set") != attempt || f.goal.Status != active || f.turn.Status != "inProgress" {
					t.Fatal("failed turn did not resume work")
				}
				if err := f.reconcile(); err != nil {
					t.Fatal(err)
				}
			}
		})
	}
}

func TestFailedCheckpoint(t *testing.T) {
	for _, status := range []goalStatus{blocked, "usageLimited"} {
		t.Run(string(status), func(t *testing.T) {
			f := newFake(t, nil)
			if err := f.reconcile(); err != nil {
				t.Fatal(err)
			}
			f.fail()
			f.goal.Status = status
			f.turn.Error.Info.Kind = "cyberPolicy"
			w := f.op.store.watch("thread-1")
			w.Goal, w.Turn, w.Armed, w.Known = f.goal, f.turn, false, true
			if status == "usageLimited" {
				w.Hold = f.goal
			}
			if err := f.op.save(); err != nil {
				t.Fatal(err)
			}
			path := f.op.store.path
			closeStore(t, f.op.store)
			s, err := openStore(path)
			if err != nil {
				t.Fatal(err)
			}
			defer closeStore(t, s)
			next := newFake(t, s)
			next.goal, next.turn = f.goal, f.turn
			if err := next.reconcile(); err != nil {
				t.Fatal(err)
			}
			want := 1
			if status == "usageLimited" {
				want = 0
			}
			if w := s.watch("thread-1"); !w.eligible() || w.Failures != want || w.Due.IsZero() {
				t.Fatal("persisted failed turn was not scheduled")
			}
		})
	}
}

func TestUsageRecovery(t *testing.T) {
	f := newFake(t, nil)
	if err := f.reconcile(); err != nil {
		t.Fatal(err)
	}
	f.fail()
	f.goal.Status = "usageLimited"
	f.turn.Error.Info.Kind = "usageLimitExceeded"
	if err := f.reconcile(); err != nil {
		t.Fatal(err)
	}
	w := f.op.store.watch("thread-1")
	if !w.eligible() || w.Failures != 0 || w.Due.IsZero() {
		t.Fatal("server usage stop was not scheduled")
	}
	allowed := true
	f.usage = &allowed
	w.Due = time.Now().Add(-time.Second)
	if err := f.reconcile(); err != nil {
		t.Fatal(err)
	}
	if f.count("thread/goal/set") != 1 || f.goal.Status != active {
		t.Fatal("server usage stop was not resumed")
	}
}

func TestFailureStop(t *testing.T) {
	for _, status := range []string{"interrupted", "paused", "blocked", "usageLimited", "cleared"} {
		t.Run(status, func(t *testing.T) {
			f := newFake(t, nil)
			if err := f.reconcile(); err != nil {
				t.Fatal(err)
			}
			f.fail()
			f.turn.Error.Info.Kind = "cyberPolicy"
			if err := f.reconcile(); err != nil {
				t.Fatal(err)
			}
			w := f.op.store.watch("thread-1")
			if !w.eligible() {
				t.Fatal("failed turn was not scheduled")
			}
			w.Due = time.Now().Add(-time.Second)
			switch status {
			case "interrupted":
				f.turn.Status = status
				f.notify("turn/completed", params{ThreadID: "thread-1", Turn: f.turn})
			case "cleared":
				appendIntent(t, f.path, intentLine("cleared"))
				f.goal = nil
			default:
				f.goal = testGoal(goalStatus(status))
				f.human(f.goal)
			}
			for range 3 {
				if err := f.reconcile(); err != nil {
					t.Fatal(err)
				}
			}
			if f.count("thread/goal/set") != 0 || w.eligible() || !w.Due.IsZero() {
				t.Fatal("explicit human interruption did not cancel recovery")
			}
			path := f.op.store.path
			closeStore(t, f.op.store)
			s, err := openStore(path)
			if err != nil {
				t.Fatal(err)
			}
			defer closeStore(t, s)
			if w := s.watch("thread-1"); w.eligible() || !w.Due.IsZero() {
				t.Fatal("restart revoked explicit human cancellation")
			}
		})
	}
}

func TestBlockSource(t *testing.T) {
	for _, source := range []string{"agent", "human"} {
		for _, status := range []string{"completed", "failed"} {
			t.Run(source+"/"+status, func(t *testing.T) {
				f := newFake(t, nil)
				if err := f.reconcile(); err != nil {
					t.Fatal(err)
				}
				f.goal, f.turn = testGoal(blocked), testTurn(status)
				if source == "human" {
					f.human(f.goal)
				} else {
					f.agent(f.goal)
				}
				f.notify("turn/completed", params{ThreadID: "thread-1", Turn: f.turn})
				if err := f.reconcile(); err != nil {
					t.Fatal(err)
				}
				w := f.op.store.watch("thread-1")
				if w.eligible() || f.count("thread/goal/set") != 0 || !w.Due.IsZero() {
					t.Fatal("explicit block did not cancel recovery")
				}
				for range 3 {
					if err := f.reconcile(); err != nil {
						t.Fatal(err)
					}
				}
				if f.count("thread/goal/set") != 0 || f.goal.Status != blocked || !w.Due.IsZero() {
					t.Fatal("explicitly blocked goal was resumed")
				}
				path := f.op.store.path
				closeStore(t, f.op.store)
				s, err := openStore(path)
				if err != nil {
					t.Fatal(err)
				}
				defer closeStore(t, s)
				next := newFake(t, s)
				next.goal, next.turn = f.goal, f.turn
				for range 2 {
					if err := next.reconcile(); err != nil {
						t.Fatal(err)
					}
				}
				if next.count("thread/goal/set") != 0 || s.watch("thread-1").eligible() {
					t.Fatal("restart revoked an explicit block")
				}
			})
		}
	}
}

func TestAgentRace(t *testing.T) {
	for _, method := range []string{"thread/goal/set", "account/rateLimits/read"} {
		t.Run(method, func(t *testing.T) {
			f := newFake(t, nil)
			if err := f.reconcile(); err != nil {
				t.Fatal(err)
			}
			f.fail()
			if method == "account/rateLimits/read" {
				f.goal.Status = "usageLimited"
				ready := true
				f.usage = &ready
			}
			seen := false
			f.before = func(m message) {
				if m.Method == method {
					seen = true
					f.before = nil
					f.goal = testGoal(blocked)
					f.agent(f.goal)
				}
			}
			if err := f.reconcile(); err != nil {
				t.Fatal(err)
			}
			w := f.op.store.watch("thread-1")
			if w.eligible() {
				w.Due = time.Now().Add(-time.Second)
			}
			for range 3 {
				if err := f.reconcile(); err != nil {
					t.Fatal(err)
				}
			}
			if !seen || w.eligible() || !w.Due.IsZero() || f.goal.Status != blocked {
				t.Fatal("in-flight recovery revoked an explicit agent block")
			}
		})
	}
}

func TestAgentResume(t *testing.T) {
	f := newFake(t, nil)
	if err := f.reconcile(); err != nil {
		t.Fatal(err)
	}
	f.fail()
	f.agent(f.goal)
	if err := f.reconcile(); err != nil {
		t.Fatal(err)
	}
	f.goal, f.turn = testGoal(active), testTurn("inProgress")
	f.turn.ID = "turn-2"
	f.human(f.goal)
	if err := f.reconcile(); err != nil {
		t.Fatal(err)
	}
	f.fail()
	f.turn.ID = "turn-2"
	if err := f.reconcile(); err != nil {
		t.Fatal(err)
	}
	w := f.op.store.watch("thread-1")
	if !w.eligible() || w.Due.IsZero() {
		t.Fatal("explicit resume did not restore server-failure recovery")
	}
	w.Due = time.Now().Add(-time.Second)
	if err := f.reconcile(); err != nil {
		t.Fatal(err)
	}
	if f.count("thread/goal/set") != 1 || f.turn.Status != "inProgress" {
		t.Fatal("server failure did not resume after explicit authorization")
	}
}