// NexusGate node agent — Go rewrite of agent/agent.js (Node.js) with 1:1
// protocol parity against the panel's /api/agent/* endpoints.
//
// Configuration comes from the process environment, exactly like the Node
// version: systemd EnvironmentFile or agent/run.sh sources
// /etc/nexusgate/agent.env before exec. This binary never reads the env
// file itself, same as agent.js.
package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"regexp"
	"strconv"
	"strings"
	"sync"
	"time"
)

// version is reported to the panel in every heartbeat. It MUST stay a plain
// semver: lib/orchestrator.js parses it with split('.').map(Number) and
// requires >= 0.6.1 before allowing AnyTLS deploys, so a suffix such as
// "0.6.8-go" would parse to NaN and break that gate.
const version = "0.6.8"

type config struct {
	Controller         string
	AgentKey           string
	SingBoxBin         string
	SingBoxStatsBin    string
	Root               string // NG_CONFIG_DIR
	ResourceDir        string // Root + "/resources"
	SingBoxConfig      string
	SingBoxVersionFile string
	SingBoxLog         string
	SingBoxCursor      string
	KeyFile            string
	SessionTokenFile   string
	TLSDir             string
	APIPort            string // NG_SINGBOX_API_PORT (string, as in JS)
	PollMS             int64
}

func envOr(name, def string) string {
	if v, ok := os.LookupEnv(name); ok {
		return v
	}
	return def
}

func loadConfig() config {
	root := envOr("NG_CONFIG_DIR", "/etc/nexusgate/agent")
	// POLL_MS = Math.max(3, Number(NG_POLL_SECONDS || 8)) * 1000.
	// JS yields NaN for a non-numeric value; we clamp to the 3s floor instead
	// (documented divergence, only affects a misconfigured env).
	pollSec := 8.0
	if raw, ok := os.LookupEnv("NG_POLL_SECONDS"); ok && raw != "" {
		if n, err := strconv.ParseFloat(raw, 64); err == nil {
			pollSec = n
		} else {
			pollSec = 0
		}
	}
	if pollSec < 3 {
		pollSec = 3
	}
	return config{
		Controller:         strings.TrimRight(os.Getenv("NG_CONTROLLER"), "/"),
		AgentKey:           os.Getenv("NG_AGENT_KEY"),
		SingBoxBin:         envOr("NG_SINGBOX_BIN", "/usr/local/bin/nexusgate-sing-box"),
		SingBoxStatsBin:    envOr("NG_SINGBOX_STATS_BIN", "/usr/local/bin/nexusgate-sing-box-stats"),
		Root:               root,
		ResourceDir:        root + "/resources",
		SingBoxConfig:      envOr("NG_SINGBOX_CONFIG", "/etc/nexusgate/sing-box/config.json"),
		SingBoxVersionFile: envOr("NG_SINGBOX_VERSION_FILE", "/etc/nexusgate/sing-box/version"),
		SingBoxLog:         envOr("NG_SINGBOX_LOG", "/var/log/nexusgate/sing-box-access.log"),
		SingBoxCursor:      envOr("NG_SINGBOX_CURSOR", "/etc/nexusgate/sing-box-cursor.json"),
		KeyFile:            envOr("NG_KEY_FILE", "/etc/nexusgate/keys.json"),
		SessionTokenFile:   envOr("NG_SESSION_TOKEN_FILE", root+"/session-token.json"),
		TLSDir:             envOr("NG_TLS_DIR", "/etc/nexusgate/tls"),
		APIPort:            envOr("NG_SINGBOX_API_PORT", "10086"),
		PollMS:             int64(pollSec * 1000),
	}
}

var cfg = loadConfig()

// stateMu serializes access to on-disk agent state (resource files, sing-box
// config, keys, cursor). agent.js gets this for free from Node's event loop
// plus synchronous fs calls; the Go goroutines need it explicitly.
var stateMu sync.Mutex

// lastEngineError holds the most recent engine failure text; serviceHealth
// surfaces it as the detail string while sing-box is down (agent.js keeps a
// module-level `lastEngineError` variable for the same purpose).
var lastEngineError struct {
	mu sync.Mutex
	s  string
}

func setLastEngineError(v string) {
	lastEngineError.mu.Lock()
	lastEngineError.s = v
	lastEngineError.mu.Unlock()
}

func getLastEngineError() string {
	lastEngineError.mu.Lock()
	defer lastEngineError.mu.Unlock()
	return lastEngineError.s
}

var httpClient = &http.Client{Timeout: 20 * time.Second}

// isoNow mirrors new Date().toISOString(): UTC, millisecond precision, Z suffix.
func isoNow() string {
	return time.Now().UTC().Format("2006-01-02T15:04:05.000Z07:00")
}

func logf(format string, args ...interface{}) {
	fmt.Printf("%s %s\n", isoNow(), fmt.Sprintf(format, args...))
}

var (
	rePrivateKey = regexp.MustCompile(`(?i)\bPrivate\s*Key\s*:\s*[^\s,;]+`)
	reBearer     = regexp.MustCompile(`(?i)\bBearer\s+[^\s,;]+`)
)

// safeError mirrors agent.js safeError(): redacts private keys and bearer
// tokens from text that may be logged or reported to the panel.
func safeError(value string) string {
	value = rePrivateKey.ReplaceAllString(value, "PrivateKey: [REDACTED]")
	value = reBearer.ReplaceAllString(value, "Bearer [REDACTED]")
	return value
}

// truncateRunes ~ JS String.slice(0, n) (JS counts UTF-16 units; close enough
// for diagnostics, and safer than cutting a UTF-8 sequence mid-rune).
func truncateRunes(s string, n int) string {
	r := []rune(s)
	if len(r) > n {
		return string(r[:n])
	}
	return s
}

func fileExists(p string) bool {
	_, err := os.Stat(p)
	return err == nil
}

// ---------- panel session token (item 25) ----------
//
// After one scrypt-verified key auth the panel issues a short-lived
// session token via the X-Agent-Session-Token response header. The
// agent persists it in its state dir (0600, same handling as
// keys.json) and presents it as the bearer on subsequent requests;
// the panel then authenticates with a fast hash compare instead of
// scrypt. Any 401 while presenting a token clears it and retries
// once with the agent key, which makes the panel issue a fresh
// token. Panels that predate session tokens never send the header,
// so the agent simply keeps authenticating with its key — the token
// layer is invisible in that direction too. This is a deliberate
// divergence from the frozen Node agent (protocol superset, wire
// compatible).

type agentSessionToken struct {
	Token     string `json:"token"`
	ExpiresAt string `json:"expiresAt"`
}

var sessionTokenState struct {
	mu     sync.Mutex
	loaded bool
	value  *agentSessionToken
}

func loadSessionTokenLocked() {
	sessionTokenState.loaded = true
	raw, err := os.ReadFile(cfg.SessionTokenFile)
	if err != nil {
		return
	}
	var tok agentSessionToken
	if err := json.Unmarshal(raw, &tok); err == nil && tok.Token != "" {
		v := tok
		sessionTokenState.value = &v
	}
}

// currentSessionToken returns the held token, or nil when none is
// stored or it is expired (a one-minute margin keeps a request from
// riding on a token that dies in flight).
func currentSessionToken() *agentSessionToken {
	sessionTokenState.mu.Lock()
	defer sessionTokenState.mu.Unlock()
	if !sessionTokenState.loaded {
		loadSessionTokenLocked()
	}
	if sessionTokenState.value == nil {
		return nil
	}
	if exp, err := time.Parse(time.RFC3339, sessionTokenState.value.ExpiresAt); err == nil && time.Until(exp) < time.Minute {
		return nil
	}
	tok := *sessionTokenState.value
	return &tok
}

func storeSessionToken(token, expiresAt string) {
	if token == "" {
		return
	}
	sessionTokenState.mu.Lock()
	defer sessionTokenState.mu.Unlock()
	if !sessionTokenState.loaded {
		loadSessionTokenLocked()
	}
	if sessionTokenState.value != nil && sessionTokenState.value.Token == token {
		return
	}
	next := &agentSessionToken{Token: token, ExpiresAt: expiresAt}
	sessionTokenState.value = next
	body, err := json.Marshal(next)
	if err != nil {
		return
	}
	tmp := cfg.SessionTokenFile + ".tmp"
	if err := os.WriteFile(tmp, append(body, '\n'), 0o600); err != nil {
		logf("Session token persist failed %s", safeError(err.Error()))
		return
	}
	if err := os.Rename(tmp, cfg.SessionTokenFile); err != nil {
		logf("Session token persist failed %s", safeError(err.Error()))
		return
	}
	_ = os.Chmod(cfg.SessionTokenFile, 0o600)
}

func clearSessionToken() {
	sessionTokenState.mu.Lock()
	defer sessionTokenState.mu.Unlock()
	sessionTokenState.value = nil
	sessionTokenState.loaded = true
	if err := os.Remove(cfg.SessionTokenFile); err != nil && !os.IsNotExist(err) {
		logf("Session token clear failed %s", safeError(err.Error()))
	}
}

// request mirrors agent.js request(): POST with bearer auth, 20s timeout,
// JSON body best-effort parsed ({} on failure), error message taken from
// body.message or synthesized from the HTTP status. Bearer selection and
// the 401 retry follow the session-token rules above.
func request(endpoint string, body []byte) (map[string]interface{}, error) {
	return requestAttempt(endpoint, body, true)
}

func requestAttempt(endpoint string, body []byte, mayRetry bool) (map[string]interface{}, error) {
	bearer := cfg.AgentKey
	usedToken := false
	if tok := currentSessionToken(); tok != nil {
		bearer = tok.Token
		usedToken = true
	}
	req, err := http.NewRequest(http.MethodPost, cfg.Controller+endpoint, bytes.NewReader(body))
	if err != nil {
		return nil, err
	}
	req.Header.Set("Authorization", "Bearer "+bearer)
	req.Header.Set("Content-Type", "application/json")
	resp, err := httpClient.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	// A freshly issued token can ride on any authenticated response.
	if issued := resp.Header.Get("X-Agent-Session-Token"); issued != "" {
		storeSessionToken(issued, resp.Header.Get("X-Agent-Session-Token-Expires"))
	}
	raw, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	if resp.StatusCode == 401 && usedToken && mayRetry {
		// Token rejected (expired or revoked server-side): drop it
		// and retry once with the key; a valid key makes the panel
		// issue a fresh token on that response.
		clearSessionToken()
		return requestAttempt(endpoint, body, false)
	}
	parsed := map[string]interface{}{}
	if len(raw) > 0 {
		_ = json.Unmarshal(raw, &parsed) // .catch(() => ({})) in JS
	}
	if resp.StatusCode < 200 || resp.StatusCode > 299 {
		if msg, ok := parsed["message"].(string); ok && msg != "" {
			return nil, fmt.Errorf("%s", msg)
		}
		return nil, fmt.Errorf("Controller returned %d", resp.StatusCode)
	}
	return parsed, nil
}

func requestJSON(endpoint string, payload interface{}) (map[string]interface{}, error) {
	body, err := marshalJSON(payload)
	if err != nil {
		return nil, err
	}
	return request(endpoint, []byte(body))
}

func ensureDirectories() error {
	if err := os.MkdirAll(cfg.ResourceDir, 0o700); err != nil {
		return err
	}
	return os.MkdirAll("/var/log/nexusgate", 0o750)
}

func sleepMS(ms int64) {
	time.Sleep(time.Duration(ms) * time.Millisecond)
}

func pollLoop() {
	for {
		result, err := request("/api/agent/poll", []byte("{}"))
		if err != nil {
			logf("Poll failed %s", err.Error())
			ms := cfg.PollMS
			if ms < 10000 {
				ms = 10000
			}
			sleepMS(ms)
			continue
		}
		jobRaw, ok := result["job"]
		if !ok || jobRaw == nil {
			sleepMS(cfg.PollMS)
			continue
		}
		encoded, _ := json.Marshal(jobRaw)
		var job agentJob
		if err := json.Unmarshal(encoded, &job); err != nil {
			logf("Poll failed %s", err.Error())
			sleepMS(cfg.PollMS)
			continue
		}
		execute(&job)
	}
}

func heartbeatLoop() {
	ticker := time.NewTicker(30 * time.Second)
	defer ticker.Stop()
	for range ticker.C {
		if err := heartbeat(); err != nil {
			logf("Heartbeat failed %s", err.Error())
		}
	}
}

func runAgent() error {
	if cfg.Controller == "" || cfg.AgentKey == "" {
		return fmt.Errorf("NG_CONTROLLER and NG_AGENT_KEY are required")
	}
	if err := ensureDirectories(); err != nil {
		return err
	}
	logf("NexusGate Agent v%s starting", version)
	// Keep the control channel alive even when a stale node configuration
	// cannot start (same rationale as agent.js main()).
	stateMu.Lock()
	if err := activateSingBoxLocked(); err != nil {
		setLastEngineError(safeError(err.Error()))
		logf("Initial sing-box activation failed %s", getLastEngineError())
	}
	stateMu.Unlock()
	if err := heartbeat(); err != nil {
		logf("Initial heartbeat failed %s", err.Error())
	}
	go heartbeatLoop()
	go usageLoop()
	pollLoop()
	return nil
}

func main() {
	// `nexusgate-agent flush-usage` mirrors `node agent.js flush-usage`,
	// invoked by agent-update.sh before the agent binary is replaced so the
	// running sing-box counters are persisted first.
	if len(os.Args) > 1 && os.Args[1] == "flush-usage" {
		if err := flushAllUsageBeforeRestart(); err != nil {
			fmt.Fprintln(os.Stderr, safeError(err.Error()))
			os.Exit(1)
		}
		return
	}
	if err := runAgent(); err != nil {
		fmt.Fprintln(os.Stderr, safeError(err.Error()))
		os.Exit(1)
	}
}