package main

import (
	"bytes"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"
	"testing"
	"time"
)

func taskFixture(t *testing.T) (*Server, http.Handler) {
	t.Helper()
	s, _, _, _ := talkFixture(t)
	var err error
	s.talkTasks, err = newTalkTaskStore(filepath.Join(t.TempDir(), "tasks.json"))
	if err != nil {
		t.Fatal(err)
	}
	s.sessions = testSession("owner")
	return s, s.talkMux()
}

func ownerTaskCall(s *Server, method, path, body string) *httptest.ResponseRecorder {
	r := httptest.NewRequest(method, "https://control.example.org"+path, strings.NewReader(body))
	r.AddCookie(&http.Cookie{Name: "wxch_session", Value: "owner"})
	r.Header.Set("Content-Type", "application/json")
	r.Header.Set("Origin", "https://control.example.org")
	w := httptest.NewRecorder()
	s.talkOwnerAPI(w, r)
	return w
}

func pairTaskWorker(t *testing.T, s *Server, h http.Handler, name string) (string, string) {
	t.Helper()
	w := ownerTaskCall(s, "POST", "/api/talk/pair", `{"name":"`+name+`"}`)
	if w.Code != 200 {
		t.Fatalf("create pairing: %d %s", w.Code, w.Body)
	}
	var pair struct{ Code string }
	json.Unmarshal(w.Body.Bytes(), &pair)
	w = talkCall(h, "POST", "/api/pair", "", `{"code":"`+pair.Code+`","workspace":"/work"}`)
	if w.Code != 200 {
		t.Fatalf("redeem: %d %s", w.Code, w.Body)
	}
	var result struct {
		Key    string
		Worker talkWorker
	}
	json.Unmarshal(w.Body.Bytes(), &result)
	if again := talkCall(h, "POST", "/api/pair", "", `{"code":"`+pair.Code+`","workspace":"/work"}`); again.Code != 401 {
		t.Fatalf("pair code must be single-use: %d", again.Code)
	}
	return result.Key, result.Worker.ID
}

func enqueueTask(t *testing.T, s *Server, workerID string) talkTask {
	t.Helper()
	w := ownerTaskCall(s, "POST", "/api/talk/tasks", `{"worker_id":"`+workerID+`","title":"CCTV","instructions":"Implement and verify CCTV"}`)
	if w.Code != 200 {
		t.Fatalf("enqueue: %d %s", w.Code, w.Body)
	}
	var task talkTask
	json.Unmarshal(w.Body.Bytes(), &task)
	return task
}

func TestTalkTasksOwnerBoundary(t *testing.T) {
	s, h := taskFixture(t)
	workerKey, workerID := pairTaskWorker(t, s, h, "x")
	peerKey, _ := s.keys.createScoped("peer", "talk")
	for _, key := range []string{"", peerKey, workerKey} {
		r := httptest.NewRequest("POST", "/api/talk/tasks", strings.NewReader(`{}`))
		r.Header.Set("Authorization", "Bearer "+key)
		w := httptest.NewRecorder()
		s.talkOwnerAPI(w, r)
		if w.Code != 401 {
			t.Fatalf("bearer cannot create owner tasks: %d", w.Code)
		}
	}
	r := httptest.NewRequest("POST", "https://control.example.org/api/talk/pair", strings.NewReader(`{"name":"evil"}`))
	r.AddCookie(&http.Cookie{Name: "wxch_session", Value: "owner"})
	r.Header.Set("Origin", "https://evil.example")
	r.Header.Set("Content-Type", "application/json")
	w := httptest.NewRecorder()
	s.talkOwnerAPI(w, r)
	if w.Code != 403 {
		t.Fatalf("foreign origin: %d", w.Code)
	}
	if _, ok := s.keys.authenticate(workerKey); ok {
		t.Fatal("receiver key opened gateway")
	}
	if w := talkCall(h, "POST", "/api/tasks/claim", peerKey, `{"wait":0}`); w.Code != 401 {
		t.Fatalf("peer key claimed work: %d", w.Code)
	}
	if w := talkCall(h, "POST", "/api/messages", workerKey, `{"text":"coordination"}`); w.Code != 200 {
		t.Fatalf("receiver cannot coordinate: %d", w.Code)
	}
	task := enqueueTask(t, s, workerID)
	if task.Source != "owner" || task.Status != "queued" {
		t.Fatalf("owner source/status: %+v", task)
	}
	if w := ownerTaskCall(s, "POST", "/api/talk/tasks", `{"worker_id":"`+workerID+`","title":"x","instructions":"x","source":"peer"}`); w.Code != 400 {
		t.Fatalf("client-supplied source accepted: %d", w.Code)
	}
}

func TestTalkWaitingAllowsOtherWorkAndIdempotentOwnerContinuation(t *testing.T) {
	s, h := taskFixture(t)
	key, id := pairTaskWorker(t, s, h, "x")
	first := enqueueTask(t, s, id)
	if w := talkCall(h, "POST", "/api/tasks/claim", key, `{"wait":0}`); w.Code != 200 {
		t.Fatal(w.Code)
	}
	body := `{"status":"waiting","note":"Host discovery required; source work saved."}`
	for i := 0; i < 2; i++ {
		if w := talkCall(h, "POST", "/api/tasks/"+first.ID, key, body); w.Code != 200 {
			t.Fatalf("waiting status delivery/retry: %d %s", w.Code, w.Body)
		}
	}
	second := enqueueTask(t, s, id)
	claimed, _, err := s.talkTasks.claim(id)
	if err != nil || claimed == nil || claimed.ID != second.ID {
		t.Fatalf("waiting task prevented independent work: %+v %v", claimed, err)
	}
	path := "/api/talk/tasks/" + first.ID + "/resume"
	var continuation talkTask
	for i := 0; i < 2; i++ {
		w := ownerTaskCall(s, "POST", path, `{}`)
		if w.Code != 200 {
			t.Fatalf("continue: %d %s", w.Code, w.Body)
		}
		var got talkTask
		json.Unmarshal(w.Body.Bytes(), &got)
		if i == 1 && got.ID != continuation.ID {
			t.Fatal("repeated continuation created another task")
		}
		continuation = got
	}
	if continuation.ResumeOf != first.ID || continuation.Status != "queued" || continuation.Note == "" || continuation.Instructions != first.Instructions {
		t.Fatalf("lost continuation context: %+v", continuation)
	}
	if s.talkTasks.data.Tasks[0].Status != "waiting" || len(s.talkTasks.data.Tasks) != 3 {
		t.Fatal("history overwritten or duplicate continuation")
	}
	if w := ownerTaskCall(s, "POST", "/api/talk/tasks/"+second.ID+"/resume", `{}`); w.Code != 409 {
		t.Fatalf("must not duplicate running work: %d", w.Code)
	}
	if next, _, _ := s.talkTasks.claim(id); next != nil {
		t.Fatal("must not overlap a running task")
	}
}

func TestTalkCancelPendingPreservesHistoryAndRefusesRunningWork(t *testing.T) {
	s, h := taskFixture(t)
	key, id := pairTaskWorker(t, s, h, "x")
	_, other := pairTaskWorker(t, s, h, "other")
	statuses := []string{"queued", "waiting", "blocked", "done", "failed"}
	for _, status := range statuses {
		enqueueTask(t, s, id)
		s.talkTasks.mu.Lock()
		i := len(s.talkTasks.data.Tasks) - 1
		s.talkTasks.data.Tasks[i].Status = status
		s.talkTasks.data.Tasks[i].Note = "Saved evidence"
		s.talkTasks.mu.Unlock()
	}
	enqueueTask(t, s, other)
	path := "/api/talk/workers/" + id + "/cancel-pending"
	r := httptest.NewRequest("POST", "https://control.example.org"+path, strings.NewReader(`{}`))
	r.Header.Set("Authorization", "Bearer "+key)
	w := httptest.NewRecorder()
	s.talkOwnerAPI(w, r)
	if w.Code != 401 {
		t.Fatal("worker credential must not cancel owner work")
	}
	s.talkTasks.data.Tasks[0].Status = "running"
	if w := ownerTaskCall(s, "POST", path, `{}`); w.Code != 409 || s.talkTasks.data.Tasks[1].Status != "waiting" {
		t.Fatal("must refuse running work atomically")
	}
	s.talkTasks.data.Tasks[0].Status = "queued"
	for attempt, count := range []string{`"cancelled":3`, `"cancelled":0`} {
		w := ownerTaskCall(s, "POST", path, `{}`)
		if w.Code != 200 || !strings.Contains(w.Body.String(), count) {
			t.Fatalf("cancel attempt%d: %d %s", attempt, w.Code, w.Body)
		}
	}
	reloaded, err := newTalkTaskStore(s.talkTasks.path)
	if err != nil || len(reloaded.data.Tasks) != 6 {
		t.Fatalf("history not preserved: %v", err)
	}
	for i, task := range reloaded.data.Tasks {
		if i < 3 && (task.Status != "cancelled" || task.Note != "Saved evidence" || task.Instructions == "") {
			t.Fatal("cancelled task lost evidence")
		}
		if i >= 3 && task.Status != append(statuses[3:], "queued")[i-3] {
			t.Fatal("cancel affected completed/failed/other-worker task")
		}
	}
	if task, _, err := s.talkTasks.claim(id); err != nil || task != nil {
		t.Fatal("cancelled tasks must not be claimed")
	}
}

func TestTalkWorkerSandboxOwnerPolicy(t *testing.T) {
	s, h := taskFixture(t)
	key, id := pairTaskWorker(t, s, h, "x")
	_, other := pairTaskWorker(t, s, h, "other")
	path := "/api/talk/workers/" + id + "/sandbox"
	if w := talkCall(h, "POST", path, key, `{"sandbox":"workspace-write"}`); w.Code == 200 {
		t.Fatal("worker must not authorize its own permissions")
	}
	if w := ownerTaskCall(s, "POST", path, `{"sandbox":"danger-full-access"}`); w.Code != 400 {
		t.Fatalf("unbounded mode accepted: %d", w.Code)
	}
	if w := ownerTaskCall(s, "POST", path, `{"sandbox":"workspace-write"}`); w.Code != 200 {
		t.Fatalf("owner policy: %d %s", w.Code, w.Body)
	}
	reloaded, err := newTalkTaskStore(s.talkTasks.path)
	if err != nil {
		t.Fatal(err)
	}
	s.talkTasks = reloaded
	enqueueTask(t, s, id)
	enqueueTask(t, s, other)
	task, _, err := s.talkTasks.claim(id)
	if err != nil || task == nil || task.Sandbox != "workspace-write" {
		t.Fatalf("claim policy: %+v %v", task, err)
	}
	second, _, err := s.talkTasks.claim(other)
	if err != nil || second == nil || second.Sandbox != "" {
		t.Fatalf("policy leaked to other worker: %+v %v", second, err)
	}
}

func TestTalkTasksClaimOwnershipAndTransitions(t *testing.T) {
	s, h := taskFixture(t)
	key, worker := pairTaskWorker(t, s, h, "first")
	otherKey, _ := pairTaskWorker(t, s, h, "second")
	task := enqueueTask(t, s, worker)
	enqueueTask(t, s, worker)
	if w := talkCall(h, "POST", "/api/tasks/"+task.ID, key, `{"status":"done"}`); w.Code != 409 {
		t.Fatalf("completed unclaimed task: %d", w.Code)
	}
	w := talkCall(h, "POST", "/api/tasks/claim", key, `{"wait":0}`)
	if w.Code != 200 || !strings.Contains(w.Body.String(), task.ID) || !strings.Contains(w.Body.String(), `"running"`) {
		t.Fatalf("claim: %d %s", w.Code, w.Body)
	}
	if w := talkCall(h, "POST", "/api/tasks/claim", key, `{"wait":0}`); !strings.Contains(w.Body.String(), `"task":null`) {
		t.Fatalf("receiver ran two jobs at once: %s", w.Body)
	}
	if w := talkCall(h, "POST", "/api/tasks/current", key, `{}`); !strings.Contains(w.Body.String(), task.ID) {
		t.Fatalf("receiver cannot reconcile its own claim: %s", w.Body)
	}
	if w := talkCall(h, "POST", "/api/tasks/current", otherKey, `{}`); !strings.Contains(w.Body.String(), `"task":null`) {
		t.Fatalf("other receiver can read the task: %s", w.Body)
	}
	if w := talkCall(h, "POST", "/api/tasks/"+task.ID, otherKey, `{"status":"done"}`); w.Code != 404 {
		t.Fatalf("other worker changed task: %d", w.Code)
	}
	for i := 0; i < 2; i++ {
		if w := talkCall(h, "POST", "/api/tasks/"+task.ID, key, `{"status":"done","note":"tested"}`); w.Code != 200 {
			t.Fatalf("terminal update/retry: %d %s", w.Code, w.Body)
		}
	}
	if w := talkCall(h, "POST", "/api/tasks/"+task.ID, key, `{"status":"running"}`); w.Code != 409 {
		t.Fatalf("reopened completed task: %d", w.Code)
	}
	if w := talkCall(h, "POST", "/api/tasks/claim", key, `{"wait":0}`); strings.Contains(w.Body.String(), `"task":null`) {
		t.Fatal("next queued task not delivered")
	}
	s.keys.setDisabled(worker, true)
	if w := talkCall(h, "POST", "/api/tasks/claim", key, `{"wait":0}`); w.Code != 401 {
		t.Fatalf("revoked receiver: %d", w.Code)
	}
	if w := ownerTaskCall(s, "GET", "/api/talk/state", ""); strings.Contains(w.Body.String(), `"hash"`) || strings.Contains(w.Body.String(), key) {
		t.Fatal("state exposed credential")
	}
}

func TestTalkTaskPersistenceAndSingleClaim(t *testing.T) {
	s, h := taskFixture(t)
	_, worker := pairTaskWorker(t, s, h, "x")
	task := enqueueTask(t, s, worker)
	var wg sync.WaitGroup
	claims := make(chan string, 8)
	for i := 0; i < 8; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			if task, _, err := s.talkTasks.claim(worker); err != nil {
				t.Error(err)
			} else if task != nil {
				claims <- task.ID
			}
		}()
	}
	wg.Wait()
	close(claims)
	if len(claims) != 1 {
		t.Fatalf("expected exactly one claim, got %d", len(claims))
	}
	restarted, err := newTalkTaskStore(s.talkTasks.path)
	if err != nil || restarted.data.Tasks[0].ID != task.ID || restarted.data.Tasks[0].Status != "running" {
		t.Fatalf("restart lost task: %v", err)
	}
	if next, _, err := restarted.claim(worker); err != nil || next != nil {
		t.Fatalf("restart must not repeat in-flight work: %+v %v", next, err)
	}
}

func TestTalkPairExpiryAndStorageFailure(t *testing.T) {
	s, h := taskFixture(t)
	s.talkTasks.pairs[hashKey(strings.Repeat("a", 24))] = talkPair{Name: "expired", Expires: time.Now().Add(-time.Second)}
	if w := talkCall(h, "POST", "/api/pair", "", `{"code":"`+strings.Repeat("a", 24)+`","workspace":"/work"}`); w.Code != 401 {
		t.Fatalf("expired code: %d", w.Code)
	}
	_, worker := pairTaskWorker(t, s, h, "x")
	s.talkTasks.path = filepath.Join(t.TempDir(), "missing", "tasks.json")
	if w := ownerTaskCall(s, "POST", "/api/talk/tasks", `{"worker_id":"`+worker+`","title":"x","instructions":"x"}`); w.Code != 500 {
		t.Fatalf("failed persistence must not acknowledge success: %d", w.Code)
	}
	if len(s.talkTasks.data.Tasks) != 0 {
		t.Fatal("failed write changed task state")
	}
	badPath := filepath.Join(t.TempDir(), "corrupt.json")
	os.WriteFile(badPath, []byte("not JSON"), 0600)
	if _, err := newTalkTaskStore(badPath); err == nil {
		t.Fatal("corrupt state must stop startup rather than discard work")
	}
}

func TestTalkReceiverEndToEndWithFakeCLI(t *testing.T) {
	python, err := exec.LookPath("python")
	if err != nil {
		t.Skip("Python 3 not installed")
	}
	s, h := taskFixture(t)
	key, worker := pairTaskWorker(t, s, h, "integration")
	task := enqueueTask(t, s, worker)
	server := httptest.NewServer(h)
	defer server.Close()
	root, err := filepath.Abs(filepath.Join("..", ".."))
	if err != nil {
		t.Fatal(err)
	}
	// The receiver talks to the real Go handlers and launches a fake local CLI, not a model.
	program := `import contextlib,importlib.util,json,pathlib,sys
c=json.load(sys.stdin)
spec=importlib.util.spec_from_file_location("receiver",c.pop("source"))
w=importlib.util.module_from_spec(spec);spec.loader.exec_module(w)
w.save_config=lambda config: None
w.receiver_lock=lambda *args: contextlib.nullcontext()
def cli(engine,result):
 return [sys.executable,"-c","import json,pathlib,sys;sys.stdin.read();pathlib.Path(sys.argv[1]).write_text(json.dumps({'status':'done','summary':'Fake CLI integration test verified.'}))",str(result)]
w.cli_command=cli
w.receive(c,once=True)
`
	input, _ := json.Marshal(map[string]any{"source": filepath.Join(root, "scripts", "talk_worker.py"), "server": server.URL, "key": key, "worker_id": worker, "name": "integration", "workspace": t.TempDir(), "engine": "codex"})
	cmd := exec.Command(python, "-c", program)
	cmd.Stdin = bytes.NewReader(input)
	if output, err := cmd.CombinedOutput(); err != nil {
		t.Fatalf("receiver integration: %v (%s)", err, output)
	}
	s.talkTasks.mu.Lock()
	defer s.talkTasks.mu.Unlock()
	got := s.talkTasks.data.Tasks[0]
	if got.ID != task.ID || got.Status != "done" || got.Note != "Fake CLI integration test verified." {
		t.Fatalf("receiver did not complete the actual task protocol: %+v", got)
	}
}