package server

import (
	"context"
	"encoding/json"
	"github.com/golang-jwt/jwt/v5"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync"
	"testing"
	"time"
	"vdi-cloud/platform/internal/model"
	"vdi-cloud/platform/internal/store"
	"vdi-cloud/platform/pkg/proxmox"
)

type batchPVE struct {
	mu             sync.Mutex
	state, outcome string
	locked, deny   bool
	writes         int
}

func fakeBatchPVE(t *testing.T, f *fixture) *batchPVE {
	p := &batchPVE{state: "stopped"}
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		p.mu.Lock()
		defer p.mu.Unlock()
		path := strings.TrimPrefix(r.URL.Path, "/api2/json")
		var data any
		switch {
		case strings.Contains(path, "/tasks/"):
			if p.outcome == "poll-failure" {
				w.WriteHeader(503)
				return
			}
			status, exit := "stopped", "OK"
			if p.outcome == "running" {
				status = "running"
			}
			if p.outcome == "failed" {
				exit = "TASK ERROR"
			}
			data = map[string]any{"status": status, "exitstatus": exit}
		case path == "/access/permissions":
			n := 1
			if p.deny {
				n = 0
			}
			data = map[string]any{r.URL.Query().Get("path"): map[string]int{"VM.PowerMgmt": n, "VM.Snapshot": n, "VM.Config.CPU": n, "VM.Config.Memory": n}}
		case strings.HasSuffix(path, "/status/current"):
			data = map[string]string{"status": p.state}
		case strings.HasSuffix(path, "/config"):
			c := map[string]any{"digest": "d1", "cores": 2, "memory": 2048, "scsi0": "local:vm-9000,size=4G"}
			if p.locked {
				c["lock"] = "backup"
			}
			data = c
		case strings.HasSuffix(path, "/pending"):
			data = []any{}
		case r.Method == "POST" && strings.Contains(path, "/status/"):
			p.writes++
			if p.outcome == "reject" {
				w.WriteHeader(403)
				return
			}
			if p.outcome == "ambiguous" {
				w.WriteHeader(500)
				return
			}
			p.state = "running"
			if strings.HasSuffix(path, "/shutdown") {
				p.state = "stopped"
			}
			data = "UPID:batch-power"
			if p.outcome == "missing-upid" {
				data = ""
			}
		default:
			t.Error("unexpected batch PVE call", r.Method, path)
			http.NotFound(w, r)
			return
		}
		_ = json.NewEncoder(w).Encode(map[string]any{"data": data})
	}))
	t.Cleanup(srv.Close)
	f.server.PVE = &proxmox.Client{BaseURL: srv.URL, HTTP: srv.Client()}
	if _, err := f.db.Exec("UPDATE desktops SET provider='proxmox',node='node-a',vmid=9000,power_state='stopped' WHERE id='desk_demo_windows'"); err != nil {
		t.Fatal(err)
	}
	return p
}
func batchRequest(action string, ids ...string) map[string]any {
	return map[string]any{"action": action, "desktopIds": ids}
}
func checkedBatch(t *testing.T, f *fixture, token string, body map[string]any) map[string]any {
	t.Helper()
	code, out := f.request("POST", "/desktop-batches/check", token, body, "")
	expect(t, 200, code, out)
	b := copySnapshotConfig(body)
	b["previewToken"] = out["previewToken"]
	return b
}
func queueBatchTest(t *testing.T, f *fixture, token string, body map[string]any, key string) map[string]any {
	t.Helper()
	code, out := f.request("POST", "/desktop-batches", token, body, key)
	expect(t, 202, code, out)
	return out
}
func batchTask(t *testing.T, f *fixture, out map[string]any) model.Task {
	t.Helper()
	for _, value := range out["items"].([]any) {
		item := value.(map[string]any)
		if id, _ := item["taskId"].(string); id != "" {
			return readSnapshotTask(t, f, id)
		}
	}
	t.Fatal("no batch task", out)
	return model.Task{}
}
func TestPowerBatchLifecycleAndPartialResults(t *testing.T) {
	for _, action := range []string{"start", "shutdown", "reboot"} {
		t.Run(action, func(t *testing.T) {
			f := newFixtureWorker(t, false)
			p := fakeBatchPVE(t, f)
			if action != "start" {
				p.state = "running"
			}
			operator := f.login("operator")
			body := checkedBatch(t, f, operator, batchRequest(action, "desk_demo_windows", "not-managed"))
			out := queueBatchTest(t, f, operator, body, "batch-once")
			batch := out["batch"].(map[string]any)
			if batch["submittedCount"] != float64(1) || batch["skippedCount"] != float64(1) || batch["activeCount"] != float64(1) || batch["successCount"] != float64(0) {
				t.Fatal(batch)
			}
			task := batchTask(t, f, out)
			f.server.execute(context.Background(), task)
			task = readSnapshotTask(t, f, task.ID)
			if task.Status != "success" || p.writes != 1 || task.ProviderRef == "" {
				t.Fatal(task, p.writes)
			}
			code, view := f.request("GET", "/desktop-batches/"+batch["id"].(string), f.login("viewer"), nil, "")
			expect(t, 200, code, view)
			counts := view["batch"].(map[string]any)
			if counts["successCount"] != float64(1) || counts["activeCount"] != float64(0) {
				t.Fatal(counts)
			}
			code, view = f.request("GET", "/desktop-batches", operator, nil, "")
			expect(t, 200, code, view)
			if view["total"] != float64(1) {
				t.Fatal(view)
			}
			// A retry after completion, even without a preview token, returns the same batch.
			delete(body, "previewToken")
			again := queueBatchTest(t, f, operator, body, "batch-once")
			if again["batch"].(map[string]any)["id"] != batch["id"] {
				t.Fatal("duplicate batch")
			}
			body["action"] = "start"
			if action == "start" {
				body["action"] = "shutdown"
			}
			code, view = f.request("POST", "/desktop-batches", operator, body, "batch-once")
			expect(t, 409, code, view)
			var submitted, completed int
			_ = f.db.Get(&submitted, "SELECT count(*) FROM audit_logs WHERE result='submitted'")
			_ = f.db.Get(&completed, "SELECT count(*) FROM audit_logs WHERE result='success' AND action=$1", action)
			if submitted != 2 || completed != 1 {
				t.Fatal("audit missing", submitted, completed)
			}
		})
	}
}
func TestPowerBatchPreviewValidationAndRBAC(t *testing.T) {
	f := newFixtureWorker(t, false)
	fakeBatchPVE(t, f)
	admin := f.login("admin")
	viewer := f.login("viewer")
	operator := f.login("operator")
	for _, path := range []string{"/desktop-batches/check", "/desktop-batches"} {
		code, out := f.request("POST", path, viewer, batchRequest("start", "desk_demo_windows"), "no")
		expect(t, 403, code, out)
	}
	for _, body := range []map[string]any{batchRequest("reset", "desk_demo_windows"), batchRequest("start"), batchRequest("start", "desk_demo_windows", "desk_demo_windows"), batchRequest("start", " desk_demo_windows"), {"action": "start", "desktopIds": []string{"desk_demo_windows"}, "force": true}} {
		code, out := f.request("POST", "/desktop-batches/check", admin, body, "")
		expect(t, 400, code, out)
	}
	ids := []string{}
	for i := 0; i < 51; i++ {
		ids = append(ids, store.ID("d_"))
	}
	code, out := f.request("POST", "/desktop-batches/check", admin, batchRequest("start", ids...), "")
	expect(t, 400, code, out)
	body := checkedBatch(t, f, admin, batchRequest("start", "desk_demo_windows"))
	preview := body["previewToken"].(string)
	code, out = f.request("GET", "/auth/me", preview, nil, "")
	expect(t, 401, code, out)
	code, out = f.request("POST", "/desktop-batches", admin, body, "")
	expect(t, 400, code, out)
	code, out = f.request("POST", "/desktop-batches", operator, body, "wrong-user")
	expect(t, 409, code, out)
	body["action"] = "shutdown"
	code, out = f.request("POST", "/desktop-batches", admin, body, "changed")
	expect(t, 409, code, out)
	body["action"] = "start"
	body["previewToken"] = admin
	code, out = f.request("POST", "/desktop-batches", admin, body, "access-not-preview")
	expect(t, 409, code, out)
	cl := &powerBatchClaims{}
	_, err := jwt.ParseWithClaims(preview, cl, func(*jwt.Token) (any, error) { return []byte(f.server.Config.JWTSecret), nil })
	if err != nil {
		t.Fatal(err)
	}
	cl.ExpiresAt = jwt.NewNumericDate(time.Now().Add(-time.Minute))
	expired, _ := jwt.NewWithClaims(jwt.SigningMethodHS256, cl).SignedString([]byte(f.server.Config.JWTSecret))
	body["previewToken"] = expired
	code, out = f.request("POST", "/desktop-batches", admin, body, "expired")
	expect(t, 409, code, out)
	body["previewToken"] = preview
	secondSession := f.login("admin")
	code, out = f.request("POST", "/desktop-batches", secondSession, body, "wrong-session")
	expect(t, 409, code, out)
	var n int
	_ = f.db.Get(&n, "SELECT count(*) FROM power_batches")
	if n != 0 {
		t.Fatal("invalid request persisted batch")
	}
}
func TestPowerBatchPreflightBlockersAndRecheck(t *testing.T) {
	f := newFixtureWorker(t, false)
	p := fakeBatchPVE(t, f)
	admin := f.login("admin")
	for _, block := range []string{"running", "lock", "deny", "active-session", "maintenance", "unknown"} {
		t.Run(block, func(t *testing.T) {
			p.state = "stopped"
			p.locked = false
			p.deny = false
			switch block {
			case "running":
				p.state = "running"
			case "lock":
				p.locked = true
			case "deny":
				p.deny = true
			case "active-session":
				_, _ = f.db.Exec("UPDATE terminals SET session_desktop_id='desk_demo_windows',last_seen=now() WHERE id='term_demo_001'")
			case "maintenance":
				_, _ = f.db.Exec("INSERT INTO tasks(id,kind,action,target_id,target_name,actor,status,idempotency_key) VALUES('block','desktop','start','desk_demo_windows','blocked','admin','queued','block')")
			case "unknown":
				_, _ = f.db.Exec("INSERT INTO tasks(id,kind,action,target_id,target_name,actor,status,idempotency_key) VALUES('block','desktop','start','desk_demo_windows','blocked','admin','timeout','block')")
			}
			code, out := f.request("POST", "/desktop-batches/check", admin, batchRequest("start", "desk_demo_windows"), "")
			expect(t, 200, code, out)
			if out["eligibleCount"] != float64(0) {
				t.Fatal(block, out)
			}
			body := batchRequest("start", "desk_demo_windows")
			body["previewToken"] = out["previewToken"]
			code, out = f.request("POST", "/desktop-batches", admin, body, "blocked-"+block)
			expect(t, 409, code, out)
			_, _ = f.db.Exec("DELETE FROM tasks WHERE id='block'")
			_, _ = f.db.Exec("UPDATE terminals SET session_desktop_id=NULL WHERE id='term_demo_001'")
		})
	}
	p.state = "stopped"
	p.locked = false
	p.deny = false
	// A desktop that was excluded from the preview cannot become a newly queued target.
	body := checkedBatch(t, f, admin, batchRequest("start", "desk_demo_windows", "desk_demo_ubuntu"))
	_, _ = f.db.Exec("UPDATE desktops SET power_state='stopped' WHERE id='desk_demo_ubuntu'")
	out := queueBatchTest(t, f, admin, body, "preview-fence")
	if out["batch"].(map[string]any)["submittedCount"] != float64(1) {
		t.Fatal("added unapproved target", out)
	}
	task := batchTask(t, f, out)
	p.state = "running"
	f.server.execute(context.Background(), task)
	if readSnapshotTask(t, f, task.ID).Status != "failed" || p.writes != 0 {
		t.Fatal("worker did not recheck before send")
	}
	// Queueing rechecks newly active sessions and records a skipped row, never sends.
	p.state = "stopped"
	body = checkedBatch(t, f, admin, batchRequest("start", "desk_demo_windows"))
	_, _ = f.db.Exec("UPDATE terminals SET session_desktop_id='desk_demo_windows',last_seen=now() WHERE id='term_demo_001'")
	out = queueBatchTest(t, f, admin, body, "session-after-preview")
	if out["batch"].(map[string]any)["submittedCount"] != float64(0) {
		t.Fatal(out)
	}
}
func TestPowerBatchUnknownAndMaintenanceProtection(t *testing.T) {
	for _, outcome := range []string{"reject", "ambiguous", "missing-upid", "poll-failure", "failed"} {
		t.Run(outcome, func(t *testing.T) {
			f := newFixtureWorker(t, false)
			p := fakeBatchPVE(t, f)
			admin := f.login("admin")
			body := checkedBatch(t, f, admin, batchRequest("start", "desk_demo_windows"))
			out := queueBatchTest(t, f, admin, body, "uncertain")
			task := batchTask(t, f, out)
			p.outcome = outcome
			f.server.execute(context.Background(), task)
			task = readSnapshotTask(t, f, task.ID)
			want := "timeout"
			if outcome == "reject" {
				want = "failed"
			}
			if task.Status != want {
				t.Fatal(task)
			}
			if want == "failed" {
				return
			}
			code, view := f.request("POST", "/desktops/desk_demo_windows/actions", admin, map[string]string{"action": "start"}, "blocked-single")
			expect(t, 409, code, view)
			_, _ = f.db.Exec("UPDATE desktops SET power_state='running' WHERE id='desk_demo_windows'")
			_, _ = f.db.Exec("UPDATE terminals SET token_hash=$1,last_seen=now() WHERE id='term_demo_001'", store.Hash("batch-device"))
			code, view = f.request("POST", "/terminal/connect-ticket", "batch-device", map[string]string{"desktopId": "desk_demo_windows"}, "")
			expect(t, 409, code, view)
			if view["code"] != "hardware_unconfirmed" {
				t.Fatal(view)
			}
			code, view = f.request("PUT", "/desktops/desk_demo_windows/connection", admin, map[string]any{"endpoint": "192.0.2.90", "width": 1920, "height": 1080, "fps": 30, "initialized": true}, "")
			expect(t, 409, code, view)
			_, _ = f.db.Exec("UPDATE tasks SET status='running' WHERE id=$1", task.ID)
			f.server.execute(context.Background(), task)
			task = readSnapshotTask(t, f, task.ID)
			if p.writes != 1 {
				t.Fatal("power replayed", p.writes)
			}
			code, view = f.request("POST", "/tasks/"+task.ID+"/batch-power-reconcile", admin, nil, "")
			expect(t, 409, code, view)
			p.outcome = ""
			f.db.SetMaxOpenConns(1)
			code, view = f.request("POST", "/tasks/"+task.ID+"/batch-power-reconcile", admin, nil, "")
			if outcome == "poll-failure" || outcome == "failed" {
				expect(t, 200, code, view)
			} else {
				expect(t, 409, code, view)
			}
			if p.writes != 1 {
				t.Fatal("read-only reconcile wrote PVE")
			}
		})
	}
}
func TestPowerBatchConcurrentReplayAndExclusivity(t *testing.T) {
	f := newFixtureWorker(t, false)
	fakeBatchPVE(t, f)
	admin := f.login("admin")
	body := checkedBatch(t, f, admin, batchRequest("start", "desk_demo_windows"))
	var wg sync.WaitGroup
	results := make(chan string, 5)
	for i := 0; i < 5; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			out := queueBatchTest(t, f, admin, body, "concurrent")
			results <- out["batch"].(map[string]any)["id"].(string)
		}()
	}
	wg.Wait()
	close(results)
	first := ""
	for id := range results {
		if first == "" {
			first = id
		}
		if id != first {
			t.Fatal("duplicate batch")
		}
	}
	code, out := f.request("POST", "/desktops/desk_demo_windows/actions", admin, map[string]string{"action": "start"}, "single-conflict")
	expect(t, 409, code, out)
	code, out = f.request("POST", "/desktops/desk_demo_windows/hardware", admin, map[string]any{"kind": "compute", "cores": 4, "sockets": 1, "memoryMiB": 4096, "digest": "d1"}, "hardware-conflict")
	expect(t, 409, code, out)
	var n int
	_ = f.db.Get(&n, "SELECT count(*) FROM tasks WHERE kind='batch-desktop'")
	if n != 1 {
		t.Fatal(n)
	}
}
func TestPowerBatchExpiryAndOrdinaryLane(t *testing.T) {
	f := newFixtureWorker(t, false)
	p := fakeBatchPVE(t, f)
	admin := f.login("admin")
	body := checkedBatch(t, f, admin, batchRequest("start", "desk_demo_windows"))
	out := queueBatchTest(t, f, admin, body, "prepared-expire")
	task := batchTask(t, f, out)
	f.server.finishBatchPower(task, "", "expired", "")
	claimed, err := f.server.claimBatchPower(context.Background(), task)
	if err != nil || claimed {
		t.Fatal("expired task claimed", err)
	}
	out = queueBatchTest(t, f, admin, body, "submitting-expire")
	task = batchTask(t, f, out)
	claimed, err = f.server.claimBatchPower(context.Background(), task)
	if err != nil || !claimed {
		t.Fatal(err)
	}
	f.server.finishBatchPower(task, "", "expired", "")
	if readSnapshotTask(t, f, task.ID).Status != "timeout" {
		t.Fatal("submitted expiry not protected")
	}
	// Clear only the isolated test job, then verify a queued batch older than 2 minutes remains eligible.
	_, _ = f.db.Exec("UPDATE power_batch_items SET stage='failed' WHERE task_id=$1", task.ID)
	_, _ = f.db.Exec("UPDATE tasks SET status='failed' WHERE id=$1", task.ID)
	out = queueBatchTest(t, f, admin, body, "long")
	task = batchTask(t, f, out)
	_, _ = f.db.Exec("UPDATE tasks SET created_at=now()-interval '3 minutes' WHERE id=$1", task.ID)
	p.outcome = "running"
	ctx, cancel := context.WithCancel(context.Background())
	done := make(chan struct{})
	go func() { f.server.RunWorker(ctx); close(done) }()
	defer func() { cancel(); <-done }()
	deadline := time.Now().Add(6 * time.Second)
	for readSnapshotTask(t, f, task.ID).ProviderRef == "" {
		if time.Now().After(deadline) {
			t.Fatal("batch failed to start")
		}
		time.Sleep(25 * time.Millisecond)
	}
	code, view := f.request("POST", "/desktops/desk_demo_kylin/actions", admin, map[string]string{"action": "start"}, "ordinary")
	expect(t, 202, code, view)
	for readSnapshotTask(t, f, view["id"].(string)).Status != "success" {
		if time.Now().After(deadline) {
			t.Fatal("batch blocked normal task lane")
		}
		time.Sleep(25 * time.Millisecond)
	}
}