package controller

import (
	"bytes"
	"net/http"
	"net/http/httptest"
	"path/filepath"
	"testing"

	"github.com/gin-gonic/gin"

	"github.com/mhsanaei/3x-ui/v3/internal/database"
	"github.com/mhsanaei/3x-ui/v3/internal/database/model"
)

// Item 7 B1: the heartbeat/register optionally carry a monitor-probe-style
// probe snapshot (facts + metrics). Reported fields land on the agent row;
// a heartbeat without a probe block (old agents) must not clear stored
// values, and a partial block only overwrites what it carries.
func TestAgentHeartbeatProbe(t *testing.T) {
	gin.SetMode(gin.TestMode)
	dbDir := t.TempDir()
	t.Setenv("XUI_DB_FOLDER", dbDir)
	if err := database.InitDB(filepath.Join(dbDir, "x-ui.db")); err != nil {
		t.Fatalf("InitDB: %v", err)
	}
	t.Cleanup(func() { _ = database.CloseDB() })
	db := database.GetDB()

	const token = "probe-test-token"
	agent := &model.AgentNode{
		TokenHash:   hashAgentToken(token),
		TokenPrefix: "probe-",
		Name:        "probe-node",
	}
	if err := db.Create(agent).Error; err != nil {
		t.Fatalf("seed agent: %v", err)
	}

	engine := gin.New()
	NewAgentController(engine.Group("/panel/api/agent"))

	post := func(path, body string) int {
		req := httptest.NewRequest(http.MethodPost, path, bytes.NewBufferString(body))
		req.Header.Set("Content-Type", "application/json")
		req.Header.Set("Authorization", "Bearer "+token)
		rec := httptest.NewRecorder()
		engine.ServeHTTP(rec, req)
		return rec.Code
	}
	load := func() model.AgentNode {
		var row model.AgentNode
		if err := db.First(&row, agent.Id).Error; err != nil {
			t.Fatalf("reload agent: %v", err)
		}
		return row
	}

	facts := `"probeFacts":{"hostname":"host-1","os":"Debian GNU/Linux 13","kernel":"6.12.1",` +
		`"arch":"amd64","virt":"kvm","cpuName":"Test CPU","cpuCores":4,` +
		`"memTotal":4000000000,"swapTotal":1000000000,"diskTotal":8000000000,` +
		`"agentVersion":"A1.0.1","ipv4":"203.0.113.9","ipv6":"2001:db8::9"}`
	metrics := `"probeMetrics":{"bootId":"boot-x/0123456789abcdef","iface":"",` +
		`"uptime":3600,"cpu":12.5,"load":[0.5,0.25,0.125],` +
		`"memTotal":4000000000,"memUsed":1500000000,"swapTotal":1000000000,"swapUsed":5000000,` +
		`"diskTotal":8000000000,"diskUsed":3000000000,` +
		`"netRxTotal":111000,"netTxTotal":222000,"netRx":1000,"netTx":2000,` +
		`"tcp":42,"udp":7,"procs":123}`

	// Register carries facts only.
	if code := post("/panel/api/agent/register",
		`{"token":"`+token+`","name":"probe-node",`+facts+`}`); code != http.StatusOK {
		t.Fatalf("register with facts: status %d", code)
	}
	row := load()
	if row.ProbeHostname != "host-1" || row.ProbeVirt != "kvm" || row.ProbeCPUCores != 4 ||
		row.ProbeMemTotal != 4000000000 || row.ProbeDiskTotal != 8000000000 {
		t.Fatalf("register facts not stored: %+v", row)
	}
	if row.ProbeAt != 0 {
		t.Fatalf("register must not stamp ProbeAt (no metrics yet): %d", row.ProbeAt)
	}

	// Heartbeat with facts+metrics stores the latest snapshot.
	if code := post("/panel/api/agent/heartbeat",
		`{"token":"`+token+`","timestamp":1,`+facts+`,`+metrics+`}`); code != http.StatusOK {
		t.Fatalf("heartbeat with probe: status %d", code)
	}
	row = load()
	if row.ProbeBootID != "boot-x/0123456789abcdef" || row.ProbeMemUsed != 1500000000 ||
		row.ProbeDiskUsed != 3000000000 || row.ProbeNetRxTot != 111000 || row.ProbeNetTxTot != 222000 ||
		row.ProbeNetRx != 1000 || row.ProbeNetTx != 2000 || row.ProbeTCP != 42 || row.ProbeUDP != 7 ||
		row.ProbeProcs != 123 || row.ProbeLoad1 != 0.5 || row.ProbeLoad15 != 0.125 || row.ProbeAt == 0 {
		t.Fatalf("heartbeat probe not stored: %+v", row)
	}
	if row.ProbeCPU < 12.4 || row.ProbeCPU > 12.6 {
		t.Fatalf("probe cpu = %v, want 12.5", row.ProbeCPU)
	}

	// Bare heartbeat (old agent): stored probe values survive untouched.
	if code := post("/panel/api/agent/heartbeat", `{"token":"`+token+`","timestamp":2}`); code != http.StatusOK {
		t.Fatalf("bare heartbeat: status %d", code)
	}
	row = load()
	if row.ProbeBootID != "boot-x/0123456789abcdef" || row.ProbeMemUsed != 1500000000 || row.ProbeHostname != "host-1" {
		t.Fatalf("bare heartbeat cleared probe: %+v", row)
	}

	// Partial metrics: only the carried field changes.
	if code := post("/panel/api/agent/heartbeat",
		`{"token":"`+token+`","timestamp":3,"probeMetrics":{"netRx":7777}}`); code != http.StatusOK {
		t.Fatalf("partial heartbeat: status %d", code)
	}
	row = load()
	if row.ProbeNetRx != 7777 || row.ProbeNetTx != 2000 || row.ProbeMemUsed != 1500000000 {
		t.Fatalf("partial probe merged wrong: %+v", row)
	}
}