// SPDX-License-Identifier: GPL-3.0-or-later

// Package baseline owns the state that makes baseline-diff detectors useful
// outside the deterministic parity fixture. It deliberately stops short of
// forensic alert snapshots; bounded envelope retention belongs to eventlog.
package baseline

import (
	"encoding/json"
	"fmt"
	"io/fs"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"sync"
	"time"

	"codeberg.org/anassaeneroi/enodia-sentinal/go-agent/internal/config"
	"codeberg.org/anassaeneroi/enodia-sentinal/go-agent/internal/model"
)

const firstSeenSchema = "enodia.first_seen.v1"

// CaptureFunc and ScanSUIDFunc keep host collection injectable. Unit tests can
// exercise the full lifecycle without reading the real /proc or filesystem.
type CaptureFunc func() (model.State, error)
type ScanSUIDFunc func() []string

type firstSeenStore struct {
	Schema             string              `json:"schema"`
	Initialized        bool                `json:"initialized"`
	PublicDestinations map[string][]string `json:"public_destinations"`
	ListenerPorts      map[string][]string `json:"listener_ports"`
}

// Manager decorates each captured State with the previous-sweep and durable
// baseline views expected by Python's four baseline detectors.
type Manager struct {
	Config   config.Config
	Now      func() time.Time
	ScanSUID ScanSUIDFunc

	mu               sync.Mutex
	initialized      bool
	started          time.Time
	lastPersistScan  time.Time
	lastSUIDScan     time.Time
	lastSUIDRefresh  time.Time
	suidScanRunning  bool
	listenerBaseline []string
	suidBaseline     []string
	suidCurrent      []string
	firstSeen        firstSeenStore
}

func New(cfg config.Config, scanSUID ScanSUIDFunc) *Manager {
	return &Manager{Config: cfg, Now: time.Now, ScanSUID: scanSUID}
}

// Initialize snapshots startup listeners and privileged files. As in Python,
// startup state becomes the baseline; stored listener/SUID files are operator
// artifacts, not inputs reloaded on daemon restart.
func (m *Manager) Initialize(capture CaptureFunc) error {
	if capture == nil {
		return fmt.Errorf("baseline capture is required")
	}
	started := m.Now()
	state, err := capture()
	if err != nil {
		return err
	}
	listeners := listenerKeys(state.Sockets)
	var suid []string
	if m.ScanSUID != nil {
		suid = sortedUnique(m.ScanSUID())
	}
	if err := os.MkdirAll(m.Config.LogDir, 0o750); err != nil {
		return err
	}
	if err := writeJSONAtomic(filepath.Join(m.Config.LogDir, "listener-baseline.json"), listeners); err != nil {
		return err
	}
	if err := writeJSONAtomic(filepath.Join(m.Config.LogDir, "suid-baseline.json"), suid); err != nil {
		return err
	}
	firstSeen := loadFirstSeen(filepath.Join(m.Config.LogDir, "first-seen.json"))

	m.mu.Lock()
	defer m.mu.Unlock()
	m.initialized = true
	m.started = started
	m.lastPersistScan = started
	m.lastSUIDRefresh = started
	m.listenerBaseline = listeners
	m.suidBaseline = suid
	m.suidCurrent = nil
	m.firstSeen = firstSeen
	return nil
}

// Prepare injects baseline state immediately before detectors run. During the
// grace window every baseline detector stays inert, matching the registry-wide
// listener_baseline gate in Python.
func (m *Manager) Prepare(state *model.State, now time.Time) error {
	if state == nil {
		return fmt.Errorf("baseline state is required")
	}
	m.mu.Lock()
	if !m.initialized {
		m.mu.Unlock()
		return fmt.Errorf("baseline manager is not initialized")
	}
	armed := now.Sub(m.started) >= m.Config.BaselineGrace
	if !armed {
		state.ListenerBaseline = nil
		state.SUIDBaseline = nil
		state.SUIDBinaries = nil
		state.PersistSince = nil
		state.PersistenceFiles = nil
		state.FirstSeenPublicDestinations = nil
		state.FirstSeenListenerPorts = nil
		m.mu.Unlock()
		return nil
	}

	state.ListenerBaseline = append([]string(nil), m.listenerBaseline...)
	state.SUIDBaseline = append([]string(nil), m.suidBaseline...)
	if m.suidCurrent != nil {
		state.SUIDBinaries = append([]string(nil), m.suidCurrent...)
	}
	since := float64(m.lastPersistScan.UnixNano()) / float64(time.Second)
	state.PersistSince = &since
	state.FirstSeenInitialized = m.firstSeen.Initialized
	state.FirstSeenPublicDestinations = cloneStringMap(m.firstSeen.PublicDestinations)
	state.FirstSeenListenerPorts = cloneStringMap(m.firstSeen.ListenerPorts)
	m.maybeStartSUIDScanLocked(now)
	m.mu.Unlock()

	// Filesystem walking can be slower than copying the small durable stores,
	// but it is bounded to explicitly configured persistence roots and remains
	// fail-open per entry.
	state.PersistenceFiles = collectPersistenceFiles(m.Config.WatchPersistence)
	return nil
}

// Commit advances the previous-sweep timestamp only after detection and saves
// first-seen additions atomically. This ordering prevents changes observed in
// the current sweep from being skipped by the next one.
func (m *Manager) Commit(state model.State, now time.Time) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	if !m.initialized {
		return fmt.Errorf("baseline manager is not initialized")
	}
	m.lastPersistScan = now
	if now.Sub(m.started) < m.Config.BaselineGrace || !m.Config.Enabled("first_seen") {
		return nil
	}
	m.firstSeen = firstSeenStore{
		Schema:             firstSeenSchema,
		Initialized:        true,
		PublicDestinations: cloneStringMap(state.FirstSeenPublicDestinations),
		ListenerPorts:      cloneStringMap(state.FirstSeenListenerPorts),
	}
	return writeJSONAtomic(filepath.Join(m.Config.LogDir, "first-seen.json"), m.firstSeen)
}

func (m *Manager) maybeStartSUIDScanLocked(now time.Time) {
	if m.ScanSUID == nil || m.suidScanRunning ||
		(!m.lastSUIDScan.IsZero() && now.Sub(m.lastSUIDScan) < m.Config.SUIDScanInterval) {
		return
	}
	m.lastSUIDScan = now
	m.suidScanRunning = true
	go func() {
		result := sortedUnique(m.ScanSUID())
		finished := m.Now()
		m.mu.Lock()
		defer m.mu.Unlock()
		m.suidCurrent = result
		m.suidScanRunning = false
		if finished.Sub(m.lastSUIDRefresh) < m.Config.SUIDRefresh {
			return
		}
		// Periodic folding matches Python: a legitimately installed privileged
		// binary eventually becomes baseline instead of alerting forever.
		m.suidBaseline = append([]string(nil), result...)
		m.lastSUIDRefresh = finished
		_ = writeJSONAtomic(filepath.Join(m.Config.LogDir, "suid-baseline.json"), result)
	}()
}

func listenerKeys(sockets []model.Socket) []string {
	keys := make([]string, 0)
	for _, socket := range sockets {
		if socket.State != "LISTEN" {
			continue
		}
		port := socket.Local
		if index := strings.LastIndex(socket.Local, ":"); index >= 0 {
			port = socket.Local[index+1:]
		}
		comm := socket.Comm
		if comm == "" {
			comm = "?"
		}
		keys = append(keys, port+"/"+comm)
	}
	return sortedUnique(keys)
}

func collectPersistenceFiles(roots []string) map[string]float64 {
	result := make(map[string]float64)
	for _, root := range roots {
		info, err := os.Lstat(root)
		if err != nil {
			continue
		}
		if !info.IsDir() {
			result[root] = float64(info.ModTime().UnixNano()) / float64(time.Second)
			continue
		}
		_ = filepath.WalkDir(root, func(path string, entry fs.DirEntry, err error) error {
			if err != nil {
				return nil
			}
			if entry.IsDir() {
				return nil
			}
			info, err := os.Lstat(path)
			if err == nil {
				result[path] = float64(info.ModTime().UnixNano()) / float64(time.Second)
			}
			return nil
		})
	}
	return result
}

func loadFirstSeen(path string) firstSeenStore {
	store := firstSeenStore{
		Schema: firstSeenSchema, PublicDestinations: map[string][]string{},
		ListenerPorts: map[string][]string{},
	}
	raw, err := os.ReadFile(path)
	if err != nil || json.Unmarshal(raw, &store) != nil {
		return firstSeenStore{
			Schema: firstSeenSchema, PublicDestinations: map[string][]string{},
			ListenerPorts: map[string][]string{},
		}
	}
	store.Schema = firstSeenSchema
	if store.PublicDestinations == nil {
		store.PublicDestinations = map[string][]string{}
	}
	if store.ListenerPorts == nil {
		store.ListenerPorts = map[string][]string{}
	}
	return store
}

func writeJSONAtomic(path string, value any) error {
	raw, err := json.MarshalIndent(value, "", "  ")
	if err != nil {
		return err
	}
	raw = append(raw, '\n')
	temporary := path + ".tmp"
	if err := os.WriteFile(temporary, raw, 0o640); err != nil {
		return err
	}
	if err := os.Rename(temporary, path); err != nil {
		return err
	}
	return os.Chmod(path, 0o640)
}

func cloneStringMap(source map[string][]string) map[string][]string {
	result := make(map[string][]string, len(source))
	for key, values := range source {
		result[key] = append([]string(nil), values...)
	}
	return result
}

func sortedUnique(values []string) []string {
	seen := make(map[string]bool, len(values))
	result := make([]string, 0, len(values))
	for _, value := range values {
		if !seen[value] {
			seen[value] = true
			result = append(result, value)
		}
	}
	sort.Strings(result)
	return result
}