package server

import (
	"archive/zip"
	"bytes"
	"encoding/json"
	"io"
	"net/http/httptest"
	"strings"
	"testing"
)

const diagnosticTestQuery = "?targetType=terminal&targetId=term_demo_001&from=2025-01-01T00:00:00Z&to=2025-01-02T00:00:00Z"

func diagnosticFixture(t *testing.T) *fixture {
	f := newFixtureWorker(t, false)
	statements := []string{
		`UPDATE terminals SET name='同名盒子' WHERE id IN ('term_demo_001','term_demo_002')`,
		`INSERT INTO tasks(id,kind,action,target_id,target_name,actor,status,message,provider_ref,idempotency_key,created_at,updated_at) VALUES
 ('diag-task','terminal','collect-logs','term_demo_001','同名盒子','admin','success','password=task-secret','provider-secret','diag-task','2025-01-01T01:00:00Z','2025-01-01T02:00:00Z'),
 ('unrelated-task','terminal','reboot','term_demo_002','同名盒子','admin','failed','unrelated-private','','unrelated','2025-01-01T01:00:00Z','2025-01-01T01:00:00Z'),
 ('gpu-task','gpu-prepare','prepare-desktop','desk_demo_windows','设计专属桌面','device:term_demo_001','timeout','gpu unknown','','gpu-task','2025-01-01T03:00:00Z','2025-01-01T03:00:00Z'),
 ('desktop-task','desktop','start','desk_demo_windows','设计专属桌面','admin','success','desktop-result','','desktop-task','2025-01-01T04:00:00Z','2025-01-01T04:00:00Z')`,
		`INSERT INTO audit_logs VALUES
 ('direct-audit','admin','update_terminal','term_demo_001','success','trace','2025-01-01T03:00:00Z'),
 ('task-audit','admin','read_diagnostics','diag-task','success','trace2','2025-01-01T03:00:00Z'),
 ('name-only-audit','admin','reboot','同名盒子','success','trace3','2025-01-01T03:00:00Z'),
 ('unrelated-audit','admin','update_terminal','term_demo_002','success','trace4','2025-01-01T03:00:00Z')`,
		`INSERT INTO session_history(id,session_id,terminal_id,terminal_name,desktop_id,desktop_name,state,source,reserved_at,connected_at,ended_at,end_reason) VALUES
 ('old-session','session-old','term_demo_001','同名盒子','desk_demo_windows','设计专属桌面','ended','ticket','2024-12-20T00:00:00Z','2024-12-20T00:01:00Z','2025-01-01T04:00:00Z','client-ended'),
 ('other-session','session-other','term_demo_002','同名盒子','desk_demo_ubuntu','研发办公桌面','ended','ticket','2025-01-01T01:00:00Z',NULL,'2025-01-01T02:00:00Z','unrelated-private')`,
		`INSERT INTO alert_rules(id,name,metric,threshold,hold_seconds) VALUES('diag-rule','设备离线','terminal_offline',120,0)`,
		`INSERT INTO alerts(id,rule_id,rule_name,metric,target_id,target_name,value,threshold,status,started_at,opened_at) VALUES
 ('term-alert','diag-rule','设备离线','terminal_offline','term_demo_001','同名盒子',121,120,'firing','2025-01-01T00:00:00Z','2025-01-01T00:00:00Z'),
 ('node-alert','diag-rule','节点异常','node_unavailable','demo-node-a','节点 A',121,120,'firing','2025-01-01T00:00:00Z','2025-01-01T00:00:00Z')`,
		`INSERT INTO task_diagnostics(task_id,content,created_at) VALUES('diag-task','{"agentVersion":"qt-test","bootId":"test-boot","kioskState":"DESKTOP_LIST","uptimeSeconds":500,"logTail":"ordinary error\nAuthorization: Bearer log-secret\npassword=private-value\nhttps://user:url-secret@example.test/path\n<script>alert(1)</script>","unexpectedPassword":"unlisted-secret"}','2025-01-01T05:00:00Z')`,
	}
	for _, statement := range statements {
		if _, err := f.db.Exec(statement); err != nil {
			t.Fatal(err)
		}
	}
	return f
}

func diagnosticDownload(f *fixture, token, query string) *httptest.ResponseRecorder {
	r := httptest.NewRequest("GET", "/api/v1/diagnostics/export"+query, nil)
	if token != "" {
		r.Header.Set("Authorization", "Bearer "+token)
	}
	w := httptest.NewRecorder()
	f.server.Router.ServeHTTP(w, r)
	return w
}
func readDiagnosticArchive(t *testing.T, body []byte) map[string][]byte {
	t.Helper()
	archive, err := zip.NewReader(bytes.NewReader(body), int64(len(body)))
	if err != nil {
		t.Fatal(err)
	}
	files := map[string][]byte{}
	for _, file := range archive.File {
		r, e := file.Open()
		if e != nil {
			t.Fatal(e)
		}
		data, e := io.ReadAll(r)
		r.Close()
		if e != nil {
			t.Fatal(e)
		}
		files[file.Name] = data
	}
	return files
}
func TestDiagnosticsStableCorrelationAndLifecycleTime(t *testing.T) {
	f := diagnosticFixture(t)
	token := f.login("operator")
	code, out := f.request("GET", "/diagnostics"+diagnosticTestQuery, token, nil, "")
	expect(t, 200, code, out)
	data, _ := json.Marshal(out)
	text := string(data)
	for _, want := range []string{"task:diag-task", "task:gpu-task", "session-ended:old-session", "audit:task-audit", "alert-open:term-alert", "log:diag-task"} {
		if !strings.Contains(text, want) {
			t.Fatal("missing correlated event", want)
		}
	}
	for _, private := range []string{"unrelated-task", "unrelated-audit", "name-only-audit", "other-session", "provider-secret", "task-secret", "log-secret", "private-value", "url-secret", "unlisted-secret"} {
		if strings.Contains(text, private) {
			t.Fatal("incorrectly included content", private)
		}
	}
	if out["counts"].(map[string]any)["logs"] != float64(1) {
		t.Fatal(out)
	}
	// Paging is performed after all sources are merged, using a stable ordering.
	code, page := f.request("GET", "/diagnostics"+diagnosticTestQuery+"&pageSize=2&page=2", token, nil, "")
	expect(t, 200, code, page)
	items := page["items"].([]any)
	all := out["items"].([]any)
	if len(items) != 2 || items[0].(map[string]any)["id"] != all[2].(map[string]any)["id"] || page["total"] != out["total"] {
		t.Fatal(page)
	}
	var tasks int
	f.db.Get(&tasks, "SELECT count(*) FROM tasks")
	if tasks != 4 {
		t.Fatal("read created tasks", tasks)
	}
}
func TestDiagnosticsDesktopScopeAndSourceFilters(t *testing.T) {
	f := diagnosticFixture(t)
	token := f.login("admin")
	query := strings.Replace(diagnosticTestQuery, "targetType=terminal&targetId=term_demo_001", "targetType=desktop&targetId=desk_demo_windows", 1)
	code, out := f.request("GET", "/diagnostics"+query, token, nil, "")
	expect(t, 200, code, out)
	data, _ := json.Marshal(out)
	text := string(data)
	if !strings.Contains(text, "alert-open:node-alert") || !strings.Contains(text, "session-ended:old-session") || strings.Contains(text, "log:diag-task") || strings.Contains(text, "other-session") {
		t.Fatal(text)
	}
	code, out = f.request("GET", "/diagnostics"+diagnosticTestQuery+"&source=logs", token, nil, "")
	expect(t, 200, code, out)
	if out["total"] != float64(1) || out["items"].([]any)[0].(map[string]any)["source"] != "logs" {
		t.Fatal(out)
	}
}
func TestDiagnosticsPermissionsAndValidation(t *testing.T) {
	f := newFixtureWorker(t, false)
	for _, role := range []string{"", "viewer", "operator", "admin"} {
		token := ""
		if role != "" {
			token = f.login(role)
		}
		want := 200
		if role == "" {
			want = 401
		}
		if role == "viewer" {
			want = 403
		}
		code, out := f.request("GET", "/diagnostics"+diagnosticTestQuery, token, nil, "")
		expect(t, want, code, out)
		w := diagnosticDownload(f, token, diagnosticTestQuery)
		if w.Code != want {
			t.Fatal(role, w.Code)
		}
	}
	token := f.login("admin")
	for _, query := range []string{"", "?targetType=all&targetId=x", diagnosticTestQuery + "&source=shell", "?targetType=terminal&targetId=term_demo_001&from=2025-01-01T00:00:00&to=2025-01-02T00:00:00Z", "?targetType=terminal&targetId=term_demo_001&from=2025-01-01T00:00:00Z&to=2025-02-01T00:00:00Z"} {
		code, out := f.request("GET", "/diagnostics"+query, token, nil, "")
		expect(t, 400, code, out)
	}
	code, out := f.request("GET", strings.Replace("/diagnostics"+diagnosticTestQuery, "term_demo_001", "missing", 1), token, nil, "")
	expect(t, 404, code, out)
}
func TestDiagnosticsZIPCompletenessAndRedaction(t *testing.T) {
	f := diagnosticFixture(t)
	token := f.login("operator")
	w := diagnosticDownload(f, token, diagnosticTestQuery+"&source=logs&page=999&pageSize=1")
	if w.Code != 200 || w.Header().Get("Content-Type") != "application/zip" || w.Header().Get("Cache-Control") != "no-store" {
		t.Fatal(w.Code, w.Body.String())
	}
	files := readDiagnosticArchive(t, w.Body.Bytes())
	if len(files) != 2 || len(files["README.txt"]) == 0 {
		t.Fatal("invalid archive entries")
	}
	var report diagnosticReport
	if err := json.Unmarshal(files["diagnostics.json"], &report); err != nil {
		t.Fatal(err)
	}
	if report.Total != 1 || len(report.Items) != 1 || report.Page != 1 || report.Items[0].Source != "logs" {
		t.Fatal(report)
	}
	text := string(files["diagnostics.json"])
	for _, secret := range []string{"provider-secret", "log-secret", "private-value", "url-secret", "unlisted-secret"} {
		if strings.Contains(text, secret) {
			t.Fatal("unredacted secret", secret)
		}
	}
	if !strings.Contains(text, "ordinary error") || !strings.Contains(text, "REDACTED") {
		t.Fatal(text)
	}
	var exports int
	f.db.Get(&exports, "SELECT count(*) FROM audit_logs WHERE action='export_diagnostics' AND target='term_demo_001' AND actor='operator'")
	if exports != 1 {
		t.Fatal("missing export audit")
	}
}
func TestDiagnosticsExportLimitAndEmptyArchive(t *testing.T) {
	f := newFixtureWorker(t, false)
	token := f.login("admin")
	if _, err := f.db.Exec(`INSERT INTO audit_logs SELECT 'diag-bulk-'||i,'admin','update_terminal','term_demo_001','success','req','2025-01-01T01:00:00Z'::timestamptz FROM generate_series(1,2001) i`); err != nil {
		t.Fatal(err)
	}
	w := diagnosticDownload(f, token, diagnosticTestQuery+"&source=audits")
	if w.Code != 413 || strings.Contains(w.Header().Get("Content-Type"), "zip") {
		t.Fatal(w.Code)
	}
	w = diagnosticDownload(f, token, diagnosticTestQuery+"&source=logs")
	if w.Code != 200 {
		t.Fatal(w.Code, w.Body.String())
	}
	files := readDiagnosticArchive(t, w.Body.Bytes())
	var report diagnosticReport
	json.Unmarshal(files["diagnostics.json"], &report)
	if report.Total != 0 || len(report.Items) != 0 {
		t.Fatal(report)
	}
}
func TestDiagnosticRedactionFormats(t *testing.T) {
	for _, value := range []string{`password="two word secret"`, `{"accessToken":"private-value"}`, "Authorization: Bearer private-value", "PVEAPIToken=admin@pve!id=private-value", "device_0123456789abcdef0123456789abcdef", "https://user:private-value@example.test/path", "-----BEGIN OPENSSH PRIVATE KEY-----\nprivate-value\n-----END OPENSSH PRIVATE KEY-----"} {
		if result := redactDiagnostic(value); strings.Contains(result, "private-value") || strings.Contains(result, "two word secret") || strings.Contains(result, "0123456789abcdef") {
			t.Fatalf("unredacted: %s", result)
		}
	}
	if redactDiagnostic("network error: offline") != "network error: offline" {
		t.Fatal("ordinary diagnostic lost")
	}
}