worker/app/models/worker_node.go

package models

import (
	"encoding/json"
	"log"
	"os"
	"strconv"
	"strings"
	"time"

	"gorm.io/datatypes"

	"rsgit.ru/rsmon/rsmon/app/models/concerns"
)

// WorkerNode represents a distributed monitoring worker
type WorkerNode struct {
	concerns.Model
	WorkerID             string         `gorm:"uniqueIndex;size:100;not null" json:"worker_id"` // UUID or configured ID
	RegionCode           string         `gorm:"size:20;not null;index" json:"region_code"`
	Region               *Region        `gorm:"foreignKey:RegionCode;references:Code" json:"region,omitempty"`
	Status               string         `gorm:"not null;default:'registered'" json:"status"` // registered, active, inactive, dead
	AuthToken            string         `gorm:"uniqueIndex;size:64;not null" json:"-"`
	LastSeen             *time.Time     `json:"last_seen"`
	Version              string         `gorm:"size:50" json:"version"`
	URL                  string         `gorm:"size:500" json:"url"`                       //nolint:lll // publicly-advertised URL; may differ from bind host:port when behind Traefik
	Capabilities         datatypes.JSON `gorm:"not null;default:'{}'" json:"capabilities"` // {"check_types": ["http","ssl","dns",...]}
	Concurrency          int            `gorm:"not null;default:20" json:"concurrency"`
	NetworkProblems      bool           `gorm:"not null;default:false;index" json:"network_problems"`
	NetworkProblemsUntil *time.Time     `json:"network_problems_until,omitempty"`
	LastFailureCount     int            `gorm:"not null;default:0" json:"last_failure_count"`
	LastTotalCount       int            `gorm:"not null;default:0" json:"last_total_count"`
	// Capability flags (see docs/plans/inventory-management.md §3).
	// All default to true so an existing worker row that pre-dates
	// this migration keeps running checks. Toggle from the admin UI
	// or POST a boolean to /api/v1/workers to disable any one of
	// them; the distworker client re-reads these flags on every
	// task poll.
	//
	// Pointer types so GORM can distinguish "client didn't supply
	// the key, fall back to the DB default" from "client explicitly
	// set false". A plain `bool` would be silently re-overwritten
	// by the column default on Save (default:true kicks in when
	// GORM sees the zero value, regardless of whether the handler
	// asked for false). See the controller tests for the
	// partial-update case.
	RunChecks      *bool `gorm:"not null;default:true" json:"run_checks"`
	CollectMetrics *bool `gorm:"not null;default:true" json:"collect_metrics"`
	DetectProjects *bool `gorm:"not null;default:true" json:"detect_projects"`
	// ServerID is the optional inventory Server this worker daemon
	// is running on (see docs/plans/servers-and-hardware-metrics.md
	// §3 and docs/plans/inventory-management.md §1). Nullable so
	// legacy "no server assigned" rows keep working. Indexed because
	// the distworker health ticker joins servers→workers frequently.
	ServerID *int64 `gorm:"type:bigint REFERENCES servers(id) ON DELETE SET NULL;index" json:"server_id,omitempty"`
	// AccountID scopes a worker to a single customer account (private
	// worker per docs/distributed/private-workers.md). NULL means a
	// platform-operated worker eligible to serve any account; non-NULL
	// is a customer-operated worker pinned to one account. ON DELETE
	// SET NULL keeps an operated worker valid if its account row is
	// ever removed without an explicit private-worker cleanup.
	AccountID *int64 `gorm:"type:bigint REFERENCES accounts(id) ON DELETE SET NULL;index" json:"account_id,omitempty"`
	LLMs      []LLM  `json:"llms,omitempty" gorm:"many2many:worker_llms;"`
	concerns.Timestamped
}

const WorkerHeartbeatFreshness = 2 * time.Minute

func (w *WorkerNode) NetworkProblemActive(now time.Time) bool {
	return w != nil && w.NetworkProblems && (w.NetworkProblemsUntil == nil || w.NetworkProblemsUntil.After(now))
}

// WorkerStatuses provides functionality.
// WorkerStatus represents the possible worker statuses
var WorkerStatuses = []string{"registered", "active", "inactive", "dead"}

// AllWorkerCheckKinds is the complete distributed-worker capability set.
func AllWorkerCheckKinds() []string {
	return []string{kindHTTP, kindSSL, kindDNS, kindSSH, kindFTP, kindWhois, kindBSSL, kindLLM, kindLLMHTTP, kindPing, kindTCP, kindUDP}
}

// NormalizeWorkerCapabilities expands aliases and removes unknown or duplicate capabilities.
func NormalizeWorkerCapabilities(capabilities []string) []string {
	allowed := make(map[string]bool)
	for _, kind := range AllWorkerCheckKinds() {
		allowed[kind] = true
	}

	seen := make(map[string]bool)
	normalized := make([]string, 0, len(capabilities))
	for _, capability := range capabilities {
		capability = strings.TrimSpace(strings.ToLower(capability))
		if capability == "all" || capability == "*" {
			return AllWorkerCheckKinds()
		}
		if !allowed[capability] || seen[capability] {
			continue
		}
		seen[capability] = true
		normalized = append(normalized, capability)
	}
	if len(normalized) == 0 {
		return AllWorkerCheckKinds()
	}
	return normalized
}

// IsAlive returns true if the worker is considered alive based on last_seen
func (w *WorkerNode) IsAlive() bool {
	if w.LastSeen == nil {
		return false
	}
	// Worker is considered dead if no heartbeat for 2 minutes
	return w.LastSeen.After(time.Now().Add(-WorkerHeartbeatFreshness))
}

// NotificationMethods returns the notification methods the worker is authorized
// to deliver (e.g. ["email", "telegram"]). Empty slice means "no notification
// delivery authorized". See docs/plans/worker-notifier-mvp.md section 4.3.
func (w *WorkerNode) NotificationMethods() []string {
	caps := w.capabilitiesMap()
	if caps == nil {
		return nil
	}
	raw, ok := caps["notification_methods"]
	if !ok {
		return nil
	}
	return parseStringList(raw)
}

// NotificationAccounts returns the account IDs the worker is authorized to
// serve for notifications. Empty slice means "owned by RSMon, all accounts".
// Customer-hosted workers (phase 4) ship a non-empty slice to scope credentials.
func (w *WorkerNode) NotificationAccounts() []int64 {
	caps := w.capabilitiesMap()
	if caps == nil {
		return nil
	}
	raw, ok := caps["notification_accounts"]
	if !ok {
		return nil
	}
	return parseInt64List(raw)
}

// AccessibleAccountIDs returns the accounts this worker may access. Empty means
// RSMon-operated/global worker. It combines the legacy single AccountID field
// with the newer notification_accounts capability list.
func (w *WorkerNode) AccessibleAccountIDs() []int64 {
	if w == nil {
		return nil
	}
	if w.AccountID != nil && *w.AccountID > 0 {
		// A private worker cannot widen its account scope through a mutable
		// capability JSON blob.
		return []int64{*w.AccountID}
	}
	seen := map[int64]bool{}
	out := []int64{}
	for _, id := range w.NotificationAccounts() {
		if id > 0 && !seen[id] {
			seen[id] = true
			out = append(out, id)
		}
	}
	return out
}

// CanDeliverNotification returns true when the worker is allowed to deliver
// the given method for the given account. An empty NotificationAccounts slice
// means the worker is RSMon-operated and may serve any account.
func (w *WorkerNode) CanDeliverNotification(method string, accountID int64) bool {
	methods := w.NotificationMethods()
	if len(methods) == 0 {
		return false
	}
	if !containsString(methods, method) {
		return false
	}
	accounts := w.AccessibleAccountIDs()
	if len(accounts) == 0 {
		return true
	}
	return containsInt64(accounts, accountID)
}

// CheckTypes returns the check kinds this worker may execute.
func (w *WorkerNode) CheckTypes() []string {
	capabilities := w.capabilitiesMap()
	if capabilities == nil {
		return nil
	}
	return parseStringList(capabilities["check_types"])
}

// SupportsTaskEnvelope is an explicit protocol capability. Version labels are
// build metadata (and may be "latest" or a commit SHA), not a wire contract.
// Rows created before this capability existed intentionally remain v1.
func (w *WorkerNode) SupportsTaskEnvelope() bool {
	capabilities := w.capabilitiesMap()
	if capabilities == nil {
		return false
	}
	supported, _ := capabilities["task_envelope"].(bool)
	return supported
}

// ReportedWorkload is the worker's local active and queued work across both
// checks and notifications. It is advisory; durable leases remain authoritative.
func (w *WorkerNode) ReportedWorkload() int {
	capabilities := w.capabilitiesMap()
	if capabilities == nil {
		return 0
	}
	keys := []string{"active_checks", "queue_depth", "active_notifications", "notification_queue_depth"}
	total := 0
	for _, key := range keys {
		switch value := capabilities[key].(type) {
		case float64:
			if value > 0 {
				total += int(value)
			}
		case int:
			if value > 0 {
				total += value
			}
		}
	}
	return total
}

func (w *WorkerNode) capabilitiesMap() map[string]interface{} {
	if w == nil || len(w.Capabilities) == 0 {
		return nil
	}
	var out map[string]interface{}
	if err := json.Unmarshal(w.Capabilities, &out); err != nil {
		return nil
	}
	return out
}

func parseStringList(raw interface{}) []string {
	switch v := raw.(type) {
	case []interface{}:
		out := make([]string, 0, len(v))
		for _, item := range v {
			if s, ok := item.(string); ok && s != "" {
				out = append(out, s)
			}
		}
		return out
	case []string:
		out := make([]string, 0, len(v))
		for _, s := range v {
			if s != "" {
				out = append(out, s)
			}
		}
		return out
	}
	return nil
}

func parseInt64List(raw interface{}) []int64 {
	switch v := raw.(type) {
	case []interface{}:
		out := make([]int64, 0, len(v))
		for _, item := range v {
			switch n := item.(type) {
			case float64:
				out = append(out, int64(n))
			case int64:
				out = append(out, n)
			}
		}
		return out
	case []int64:
		return v
	}
	return nil
}

func containsString(haystack []string, needle string) bool {
	for _, s := range haystack {
		if s == needle {
			return true
		}
	}
	return false
}

func containsInt64(haystack []int64, needle int64) bool {
	for _, n := range haystack {
		if n == needle {
			return true
		}
	}
	return false
}

// EnsureConfiguredWorkerNode creates or updates the bundled Docker Compose worker from environment variables.
func EnsureConfiguredWorkerNode() {
	token := envFirst("WORKER_AUTH_TOKEN", "RSMON_AUTH_TOKEN")
	if token == "" {
		return
	}

	deployEnv := envDefault("DEPLOY_ENV", "production")
	workerID := envDefault("RSMON_WORKER_ID", "worker-"+deployEnv+"-01")
	regionCode := envDefault("RSMON_REGION_CODE", deployEnv)
	version := envDefault("RSMON_WORKER_VERSION", envDefault("IMAGE_TAG", "latest"))
	concurrency := envInt("WORKER_CONCURRENCY", 20)
	workerURL := strings.TrimSpace(os.Getenv("WORKER_URL"))
	capabilities := NormalizeWorkerCapabilities(splitEnvList(envDefault(
		"RSMON_CAPABILITIES",
		"http,ssl,dns,ssh,ftp,whois,bssl,llm,llm-http,ping,tcp,udp",
	)))

	region := Region{}
	if err := DB().Where("code = ?", regionCode).First(&region).Error; err != nil {
		region = Region{Code: regionCode, Name: regionCode, Enabled: true}
		if err := DB().Create(&region).Error; err != nil {
			log.Printf("worker: failed to create configured worker region %s: %v", regionCode, err)
			return
		}
	}

	capJSON, err := json.Marshal(map[string]interface{}{
		"check_types":           capabilities,
		"task_envelope":         true,
		"notification_methods":  parseStringList(notificationMethodsFromEnv()),
		"notification_accounts": parseInt64List(notificationAccountsFromEnv()),
	})
	if err != nil {
		log.Printf("worker: failed to marshal configured worker capabilities: %v", err)
		return
	}

	worker := WorkerNode{}
	DB().Where("worker_id = ? OR auth_token = ?", workerID, token).First(&worker)
	created := worker.ID == 0
	worker.WorkerID = workerID
	worker.RegionCode = regionCode
	worker.Status = "registered"
	worker.AuthToken = token
	worker.Version = version
	worker.URL = workerURL
	worker.Capabilities = datatypes.JSON(capJSON)
	worker.Concurrency = concurrency

	if err := DB().Save(&worker).Error; err != nil {
		log.Printf("worker: failed to provision configured worker %s: %v", workerID, err)
		return
	}
	if created {
		log.Printf("worker: provisioned configured worker %s in region %s", workerID, regionCode)
	} else {
		log.Printf("worker: updated configured worker %s in region %s", workerID, regionCode)
	}
}

func splitEnvList(value string) []string {
	parts := strings.Split(value, ",")
	items := make([]string, 0, len(parts))
	for _, part := range parts {
		part = strings.TrimSpace(part)
		if part != "" {
			items = append(items, part)
		}
	}
	return items
}

func envFirst(keys ...string) string {
	for _, key := range keys {
		if value := os.Getenv(key); value != "" {
			return value
		}
	}
	return ""
}

func envDefault(key, defaultValue string) string {
	if value := os.Getenv(key); value != "" {
		return value
	}
	return defaultValue
}

func envInt(key string, defaultValue int) int {
	value, err := strconv.Atoi(os.Getenv(key))
	if err != nil || value <= 0 {
		return defaultValue
	}
	return value
}

// Notification method constants used across the package. Centralized here so
// the literal does not appear three or more times (goconst).
const (
	methodEmail      = "email"
	methodTelegram   = "telegram"
	methodWebhook    = "webhook"
	methodMattermost = "mattermost"
)

// defaultNotificationMethods is the operated-worker default notification method
// list. Lives at package scope so goconst does not flag the literal across
// the package (account.go and user.go already reference "email").
var defaultNotificationMethods = []string{methodEmail, methodTelegram, methodWebhook, methodMattermost}

// notificationMethodsFromEnv reads the optional NOTIFICATION_METHODS env var.
// Empty result yields the default "all four" list so the operated worker can
// deliver email / telegram / webhook / mattermost out of the box.
func notificationMethodsFromEnv() []string {
	raw := strings.TrimSpace(os.Getenv("NOTIFICATION_METHODS"))
	if raw == "" {
		return append([]string{}, defaultNotificationMethods...)
	}
	out := make([]string, 0, 4)
	for _, part := range strings.Split(raw, ",") {
		part = strings.TrimSpace(part)
		if part != "" {
			out = append(out, part)
		}
	}
	return out
}

// notificationAccountsFromEnv reads the optional NOTIFICATION_ACCOUNTS env var.
// Empty result means "all accounts allowed" (the RSMon-operated default).
func notificationAccountsFromEnv() []int64 {
	raw := strings.TrimSpace(os.Getenv("NOTIFICATION_ACCOUNTS"))
	if raw == "" {
		return nil
	}
	out := make([]int64, 0, 4)
	for _, part := range strings.Split(raw, ",") {
		part = strings.TrimSpace(part)
		if part == "" {
			continue
		}
		n, err := strconv.ParseInt(part, 10, 64)
		if err != nil || n <= 0 {
			continue
		}
		out = append(out, n)
	}
	return out
}