package server

import (
	"bytes"
	"context"
	"encoding/json"
	"encoding/pem"
	"fmt"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"vdi-cloud/platform/internal/config"
	"vdi-cloud/platform/internal/model"
	"vdi-cloud/platform/internal/store"
	"vdi-cloud/platform/pkg/proxmox"
)

func TestPVEConnectionEncryptedReloadAndRestrictions(t *testing.T) {
	f := newFixtureWorker(t, false)
	admin := f.login("admin")
	var calls atomic.Int32
	srv := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		calls.Add(1)
		if r.Method != "GET" || r.URL.Path != "/api2/json/cluster/resources" {
			t.Error("unexpected connection write", r.Method, r.URL.Path)
		}
		if r.Header.Get("Authorization") != "PVEAPIToken=test@pve!console=private-test-token" {
			w.WriteHeader(401)
			return
		}
		_ = json.NewEncoder(w).Encode(map[string]any{"data": []map[string]any{{"type": "node", "node": "test-node", "status": "online"}}})
	}))
	defer srv.Close()
	ca := string(pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: srv.Certificate().Raw}))
	f.server.Config.PVEURL = srv.URL
	f.server.Config.PVETokenID = "test@pve!console"
	f.server.Config.PVETokenSecret = "private-test-token"
	code, out := f.request("GET", "/settings/pve", admin, nil, "")
	expect(t, 200, code, out)
	if out["source"] != "environment" || out["tokenConfigured"] != true || out["tokenSecret"] != nil {
		t.Fatal(out)
	}
	body := map[string]any{"url": srv.URL, "tokenId": "test@pve!console", "tokenSecret": "", "caPem": ca, "revision": "environment"}
	bad := map[string]any{}
	for k, v := range body {
		bad[k] = v
	}
	bad["caPem"] = ""
	code, out = f.request("POST", "/settings/pve/check", admin, bad, "")
	expect(t, 502, code, out)
	code, out = f.request("POST", "/settings/pve/check", admin, body, "")
	expect(t, 200, code, out)
	code, out = f.request("PUT", "/settings/pve", f.login("operator"), body, "")
	expect(t, 403, code, out)
	code, out = f.request("PUT", "/settings/pve", admin, body, "")
	expect(t, 200, code, out)
	var encrypted string
	if err := f.db.Get(&encrypted, "SELECT encrypted_config FROM pve_connection"); err != nil {
		t.Fatal(err)
	}
	if strings.Contains(encrypted, "private-test-token") || strings.Contains(encrypted, "test@pve") {
		t.Fatal("plaintext credential persisted")
	}
	code, out = f.request("GET", "/settings/pve", admin, nil, "")
	expect(t, 200, code, out)
	raw, _ := json.Marshal(out)
	if bytes.Contains(raw, []byte("private-test-token")) || out["source"] != "database" {
		t.Fatal("credential leaked", out)
	}
	revision := out["revision"]
	replica := New(f.server.Config, f.db, nil)
	if err := replica.RefreshCompute(context.Background()); err != nil {
		t.Fatal(err)
	}
	if replica.currentPVE() == nil || replica.currentPVE().TokenSecret != "private-test-token" {
		t.Fatal("replica failed to load")
	}
	code, out = f.request("PUT", "/settings/pve", admin, body, "")
	expect(t, 409, code, out)
	body["revision"] = revision
	body["tokenId"] = "other@pve!console"
	code, out = f.request("PUT", "/settings/pve", admin, body, "")
	expect(t, 400, code, out)
	body["tokenSecret"] = "wrong"
	code, out = f.request("PUT", "/settings/pve", admin, body, "")
	expect(t, 502, code, out)
	body["url"] = "https://different.example:8006"
	code, out = f.request("PUT", "/settings/pve", admin, body, "")
	expect(t, 409, code, out)
	body["url"] = srv.URL
	body["tokenId"] = "test@pve!console"
	body["tokenSecret"] = ""
	code, out = f.request("PUT", "/settings/pve", admin, body, "")
	expect(t, 200, code, out)
	if err := replica.RefreshCompute(context.Background()); err != nil {
		t.Fatal(err)
	}
	if replica.computeRevision == revision {
		t.Fatal("replica kept stale revision")
	}
	replica.Config.JWTSecret = strings.Repeat("different", 10)
	replica.computeRevision = ""
	if err := replica.RefreshCompute(context.Background()); err == nil {
		t.Fatal("wrong encryption key accepted")
	}
	var leaked bool
	f.db.Get(&leaked, "SELECT EXISTS(SELECT 1 FROM audit_logs WHERE target LIKE '%private-test-token%' OR action LIKE '%private-test-token%')")
	if leaked || calls.Load() < 3 {
		t.Fatal("audit leak or connection never checked")
	}
}

func loginRecorder(s *Server, username, remote, forwarded string) *httptest.ResponseRecorder {
	raw, _ := json.Marshal(map[string]string{"username": username, "password": "wrong-password"})
	r := httptest.NewRequest("POST", "/api/v1/auth/login", bytes.NewReader(raw))
	r.RemoteAddr = remote
	r.Header.Set("Content-Type", "application/json")
	r.Header.Set("X-Forwarded-For", forwarded)
	w := httptest.NewRecorder()
	s.Router.ServeHTTP(w, r)
	return w
}
func TestLoginLimitsPersistentConcurrentAndForwarding(t *testing.T) {
	f := newFixtureWorker(t, false)
	for i := 0; i < 10; i++ {
		w := loginRecorder(f.server, "missing", "192.0.2.1:1000", fmt.Sprintf("198.51.100.%d", i))
		if w.Code != 401 {
			t.Fatal(w.Code, w.Body.String())
		}
	}
	replica := New(f.server.Config, f.db, nil)
	w := loginRecorder(replica, " MISSING ", "192.0.2.2:1000", "")
	if w.Code != 429 || w.Header().Get("Retry-After") == "" {
		t.Fatal("account budget reset", w.Code)
	}
	if _, err := f.db.Exec("UPDATE login_limits SET window_started=now()-interval '6 minutes'"); err != nil {
		t.Fatal(err)
	}
	w = loginRecorder(replica, "missing", "192.0.2.1:1000", "")
	if w.Code != 401 {
		t.Fatal("window did not reopen")
	}
	// Each account is different; spoofed forwarding must not bypass the source budget.
	for i := 0; i < 29; i++ {
		w = loginRecorder(f.server, fmt.Sprintf("account-%d", i), "192.0.2.1:1000", fmt.Sprintf("203.0.113.%d", i))
		if w.Code != 401 {
			t.Fatal(i, w.Code)
		}
	}
	w = loginRecorder(f.server, "last-account", "192.0.2.1:1000", "203.0.113.250")
	if w.Code != 429 {
		t.Fatal("forged forwarding bypassed limit", w.Code)
	}
	var wg sync.WaitGroup
	var limited atomic.Int32
	for i := 0; i < 20; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			if loginRecorder(f.server, "parallel", "192.0.2.9:1000", "").Code == 429 {
				limited.Add(1)
			}
		}()
	}
	wg.Wait()
	if limited.Load() != 10 {
		t.Fatal("non-atomic account budget", limited.Load())
	}
	var key string
	if err := f.db.Get(&key, "SELECT key FROM login_limits LIMIT 1"); err != nil || len(key) != 64 {
		t.Fatal("cleartext login key")
	}
	// Only the explicitly trusted proxy may supply a client address.
	cfg := config.Config{Mode: "demo", JWTSecret: f.server.Config.JWTSecret, TrustedProxies: []string{"10.20.0.2"}}
	proxy := New(cfg, f.db, nil)
	for i := 0; i < 31; i++ {
		w = loginRecorder(proxy, fmt.Sprintf("proxy-%d", i), "10.20.0.2:80", "198.51.100.10")
		if i == 30 && w.Code != 429 {
			t.Fatal("trusted proxy IP budget ignored")
		}
	}
}

type hardwarePVE struct {
	mu       sync.Mutex
	config   map[string]any
	state    string
	pending  bool
	deny     bool
	lowSpace bool
	outcome  string
	posts    int
	forms    []map[string]string
}

func fakeHardwarePVE(t *testing.T, f *fixture) *hardwarePVE {
	p := &hardwarePVE{state: "stopped", config: map[string]any{"digest": "d1", "cores": 2, "sockets": 1, "memory": 2048, "balloon": 1024, "vcpus": 2, "scsi0": "local:vm-9000,size=4G", "hostpci0": "0000:01:00.0", "cpu": "host", "args": "preserved"}}
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		p.mu.Lock()
		defer p.mu.Unlock()
		var data any
		path := strings.TrimPrefix(r.URL.Path, "/api2/json")
		switch {
		case strings.Contains(path, "/tasks/"):
			if p.outcome == "poll-failure" {
				w.WriteHeader(503)
				return
			}
			data = map[string]any{"status": "stopped", "exitstatus": "OK"}
		case path == "/access/permissions":
			yes := 1
			if p.deny {
				yes = 0
			}
			data = map[string]any{r.URL.Query().Get("path"): map[string]int{"VM.Config.CPU": yes, "VM.Config.Memory": yes, "VM.Config.Disk": yes, "Datastore.AllocateSpace": yes}}
		case path == "/nodes/node-a/storage":
			space := float64(100 << 30)
			if p.lowSpace {
				space = 0
			}
			data = []map[string]any{{"storage": "local", "avail": space, "active": 1, "enabled": 1}}
		case strings.HasSuffix(path, "/status/current"):
			data = map[string]any{"status": p.state}
		case strings.HasSuffix(path, "/pending"):
			data = []map[string]any{}
			if p.pending {
				data = []map[string]any{{"key": "memory", "pending": 4096}}
			}
		case strings.HasSuffix(path, "/config") && r.Method == "GET":
			if r.URL.Query().Get("current") != "1" {
				t.Error("read pending configuration as current")
			}
			data = p.config
		case strings.HasSuffix(path, "/config") || strings.HasSuffix(path, "/resize"):
			p.posts++
			_ = r.ParseForm()
			form := map[string]string{}
			for k := range r.Form {
				form[k] = r.Form.Get(k)
			}
			p.forms = append(p.forms, form)
			if p.outcome == "reject" {
				w.WriteHeader(403)
				return
			}
			if p.outcome == "ambiguous" {
				w.WriteHeader(500)
				return
			}
			if strings.HasSuffix(path, "/resize") {
				if r.Method != "PUT" || form["disk"] != "scsi0" || form["size"] != "8G" {
					t.Error("unsafe disk resize", r.Method, form)
				}
				p.config["scsi0"] = "local:vm-9000,size=8G"
			} else {
				if r.Method != "POST" {
					t.Error("wrong config method")
				}
				for _, key := range []string{"cores", "sockets", "memory", "vcpus"} {
					if v := form[key]; v != "" {
						var n int
						fmt.Sscan(v, &n)
						p.config[key] = n
					}
				}
				for key := range form {
					if !strings.Contains("|cores|sockets|memory|vcpus|digest|", "|"+key+"|") {
						t.Error("non-whitelisted config write", key)
					}
				}
			}
			p.config["digest"] = "d2"
			data = "UPID:hardware-test"
		default:
			t.Error("unexpected hardware request", 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 runHardwareTask(t *testing.T, f *fixture, admin string, body map[string]any, key string) model.Task {
	code, out := f.request("POST", "/desktops/desk_demo_windows/hardware", admin, body, key)
	expect(t, 202, code, out)
	var task model.Task
	if err := f.db.Get(&task, "SELECT * FROM tasks WHERE id=$1", out["id"]); err != nil {
		t.Fatal(err)
	}
	f.server.execute(context.Background(), task)
	if err := f.db.Get(&task, "SELECT * FROM tasks WHERE id=$1", task.ID); err != nil {
		t.Fatal(err)
	}
	return task
}
func TestHardwarePreflightTasksAndDiskAbsoluteSize(t *testing.T) {
	f := newFixtureWorker(t, false)
	p := fakeHardwarePVE(t, f)
	admin := f.login("admin")
	body := map[string]any{"kind": "compute", "cores": 4, "sockets": 1, "memoryMiB": 4096, "digest": "d1"}
	endpoint := "/desktops/desk_demo_windows/hardware"
	code, out := f.request("GET", endpoint, f.login("viewer"), nil, "")
	expect(t, 200, code, out)
	code, out = f.request("POST", endpoint+"/check", f.login("operator"), body, "")
	expect(t, 403, code, out)
	for _, block := range []string{"running", "pending", "lock", "balloon", "digest", "deny"} {
		p.mu.Lock()
		switch block {
		case "running":
			p.state = "running"
		case "pending":
			p.pending = true
		case "lock":
			p.config["lock"] = "backup"
		case "balloon":
			p.config["balloon"] = 8192
		case "digest":
			p.config["digest"] = "changed"
		case "deny":
			p.deny = true
		}
		p.mu.Unlock()
		code, out = f.request("POST", endpoint+"/check", admin, body, "")
		expect(t, 409, code, out)
		p.mu.Lock()
		p.state = "stopped"
		p.pending = false
		delete(p.config, "lock")
		p.config["balloon"] = 1024
		p.config["digest"] = "d1"
		p.deny = false
		p.mu.Unlock()
	}
	body["cpu"] = "malicious"
	code, out = f.request("POST", endpoint, admin, body, "unknown-field")
	expect(t, 400, code, out)
	delete(body, "cpu")
	task := runHardwareTask(t, f, admin, body, "compute-once")
	if task.Status != "success" {
		t.Fatal(task)
	}
	var submitted int
	if err := f.db.Get(&submitted, "SELECT count(*) FROM audit_logs WHERE action='configure-hardware' AND result='submitted'"); err != nil || submitted != 1 {
		t.Fatal("submission audit is not distinct from completion", submitted, err)
	}
	code, out = f.request("POST", endpoint, admin, body, "compute-once")
	expect(t, 202, code, out)
	if out["id"] != task.ID {
		t.Fatal("replay lost task")
	}
	body["memoryMiB"] = 8192
	code, out = f.request("POST", endpoint, admin, body, "compute-once")
	expect(t, 409, code, out)
	p.mu.Lock()
	if p.posts != 1 || p.config["hostpci0"] != "0000:01:00.0" || p.config["cpu"] != "host" || p.config["args"] != "preserved" {
		t.Fatal("hardware changed twice or passthrough altered")
	}
	p.mu.Unlock()
	disk := map[string]any{"kind": "disk", "disk": "scsi0", "sizeGiB": 8, "digest": "d2"}
	p.mu.Lock()
	p.lowSpace = true
	p.mu.Unlock()
	code, out = f.request("POST", endpoint+"/check", admin, disk, "")
	expect(t, 409, code, out)
	p.mu.Lock()
	p.lowSpace = false
	p.mu.Unlock()
	task = runHardwareTask(t, f, admin, disk, "disk-once")
	if task.Status != "success" {
		t.Fatal(task)
	}
	code, out = f.request("POST", endpoint+"/check", admin, disk, "")
	expect(t, 409, code, out)
}
func TestHardwareUnknownRestartReconcileAndExpiration(t *testing.T) {
	for _, outcome := range []string{"ambiguous", "poll-failure", "reject"} {
		t.Run(outcome, func(t *testing.T) {
			f := newFixtureWorker(t, false)
			p := fakeHardwarePVE(t, f)
			p.outcome = outcome
			admin := f.login("admin")
			body := map[string]any{"kind": "compute", "cores": 4, "sockets": 1, "memoryMiB": 4096, "digest": "d1"}
			task := runHardwareTask(t, f, admin, body, "unknown")
			want := "timeout"
			if outcome == "reject" {
				want = "failed"
			}
			if task.Status != want {
				t.Fatal(task)
			}
			if want == "failed" {
				return
			}
			code, out := f.request("POST", "/desktops/desk_demo_windows/actions", admin, map[string]string{"action": "start"}, "power")
			expect(t, 409, code, out)
			code, out = f.request("POST", "/desktops/desk_demo_windows/hardware", admin, body, "another")
			expect(t, 409, code, out)
			// A worker restarted after submitting must never issue the mutation again.
			if _, err := f.db.Exec("UPDATE tasks SET status='running' WHERE id=$1", task.ID); err != nil {
				t.Fatal(err)
			}
			f.server.execute(context.Background(), task)
			p.mu.Lock()
			if p.posts != 1 {
				t.Fatal("hardware replayed")
			}
			p.outcome = ""
			p.mu.Unlock()
			code, out = f.request("POST", "/tasks/"+task.ID+"/hardware-reconcile", admin, nil, "")
			if outcome == "poll-failure" {
				expect(t, 200, code, out)
			} else {
				expect(t, 409, code, out)
			}
		})
	}
	f := newFixtureWorker(t, false)
	fakeHardwarePVE(t, f)
	admin := f.login("admin")
	body := map[string]any{"kind": "compute", "cores": 4, "sockets": 1, "memoryMiB": 4096, "digest": "d1"}
	code, out := f.request("POST", "/desktops/desk_demo_windows/hardware", admin, body, "expire")
	expect(t, 202, code, out)
	var task model.Task
	_ = f.db.Get(&task, "SELECT * FROM tasks WHERE id=$1", out["id"])
	f.server.finishHardware(task, "", "expired", "")
	var stage string
	_ = f.db.Get(&stage, "SELECT stage FROM hardware_jobs WHERE task_id=$1", task.ID)
	if stage != "failed" {
		t.Fatal(stage)
	}
	claimed, err := f.server.claimHardware(context.Background(), task)
	if err != nil || claimed {
		t.Fatal("expired task could submit", err)
	}
	// Power and specification submissions share a target lock and unique active slot.
	code, out = f.request("POST", "/desktops/desk_demo_windows/hardware", admin, body, "submit-expire")
	expect(t, 202, code, out)
	if err = f.db.Get(&task, "SELECT * FROM tasks WHERE id=$1", out["id"]); err != nil {
		t.Fatal(err)
	}
	code, out = f.request("POST", "/desktops/desk_demo_windows/actions", admin, map[string]string{"action": "start"}, "blocked-power")
	expect(t, 409, code, out)
	claimed, err = f.server.claimHardware(context.Background(), task)
	if err != nil || !claimed {
		t.Fatal("could not claim prepared task", err)
	}
	f.server.finishHardware(task, "", "expired", "")
	if err = f.db.Get(&stage, "SELECT stage FROM hardware_jobs WHERE task_id=$1", task.ID); err != nil || stage != "uncertain" {
		t.Fatal(stage, err)
	}
	// Even if inventory still says running, uncertain specifications block new tickets.
	if _, err = f.db.Exec("UPDATE desktops SET power_state='running' WHERE id='desk_demo_windows'"); err != nil {
		t.Fatal(err)
	}
	if _, err = f.db.Exec("UPDATE terminals SET token_hash=$1,last_seen=now() WHERE id='term_demo_001'", store.Hash("hardware-device")); err != nil {
		t.Fatal(err)
	}
	code, out = f.request("POST", "/terminal/connect-ticket", "hardware-device", map[string]string{"desktopId": "desk_demo_windows"}, "")
	expect(t, 409, code, out)
	if out["code"] != "hardware_unconfirmed" {
		t.Fatal("ticket did not check hardware maintenance", out)
	}
}