package main

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"net/url"
	"regexp"
	"sync"
	"time"
)

var bearerSecret = regexp.MustCompile(`(?i)(authorization\s*[:=]\s*bearer\s+)[^\s,;&]+`)
var namedSecret = regexp.MustCompile(`(?i)((?:password|passwd|token|secret|api[_-]?key)["']?\s*[:=]\s*)(?:"[^"]*"|'[^']*'|[^\s,;&]+)`)

type Poller struct {
	store           *Store
	agent           *AgentClient
	metricRetention time.Duration
	logRetention    time.Duration
}

func (poller *Poller) Run(ctx context.Context) {
	pollTicker := time.NewTicker(time.Minute)
	cleanupTicker := time.NewTicker(time.Hour)
	defer pollTicker.Stop()
	defer cleanupTicker.Stop()
	poller.poll(ctx)
	for {
		select {
		case <-ctx.Done():
			return
		case <-pollTicker.C:
			poller.poll(ctx)
		case <-cleanupTicker.C:
			if _, err := poller.store.pool.Exec(ctx, "DELETE FROM metric_samples WHERE sampled_at < $1", time.Now().Add(-poller.metricRetention)); err != nil {
				log.Printf("metric retention cleanup: %v", err)
			}
			if _, err := poller.store.pool.Exec(ctx, "DELETE FROM agent_log_entries WHERE event_at < $1", time.Now().Add(-poller.logRetention)); err != nil {
				log.Printf("log retention cleanup: %v", err)
			}
			_, _ = poller.store.pool.Exec(ctx, "DELETE FROM sessions WHERE expires_at < now()")
		}
	}
}

func (poller *Poller) poll(ctx context.Context) {
	servers, err := poller.store.Servers(ctx)
	if err != nil {
		log.Printf("poll server list: %v", err)
		return
	}
	jobs := make(chan Server)
	var wait sync.WaitGroup
	workerCount := 4
	if len(servers) < workerCount {
		workerCount = len(servers)
	}
	for worker := 0; worker < workerCount; worker++ {
		wait.Add(1)
		go func() {
			defer wait.Done()
			for server := range jobs {
				poller.pollServer(ctx, server)
			}
		}()
	}
	for _, server := range servers {
		select {
		case jobs <- server:
		case <-ctx.Done():
			close(jobs)
			wait.Wait()
			return
		}
	}
	close(jobs)
	wait.Wait()
}

func (poller *Poller) pollServer(ctx context.Context, server Server) {
	statusCode, body, err := poller.agent.Do(ctx, server, http.MethodGet, "/v1/system/status", nil)
	if err != nil || statusCode != http.StatusOK {
		_, _ = poller.store.pool.Exec(ctx, "UPDATE servers SET status='offline',updated_at=now() WHERE id=$1", server.ID)
		return
	}
	var status struct {
		Hostname    string  `json:"hostname"`
		CPU         float64 `json:"cpu_usage_percent"`
		MemoryTotal uint64  `json:"memory_total_bytes"`
		MemoryUsed  uint64  `json:"memory_used_bytes"`
		Filesystems []struct {
			Total uint64 `json:"total_bytes"`
			Used  uint64 `json:"used_bytes"`
		} `json:"filesystems"`
	}
	if json.Unmarshal(body, &status) != nil {
		return
	}
	memory := percent(status.MemoryUsed, status.MemoryTotal)
	var diskTotal, diskUsed uint64
	for _, filesystem := range status.Filesystems {
		diskTotal += filesystem.Total
		diskUsed += filesystem.Used
	}
	disk := percent(diskUsed, diskTotal)
	var rx, tx int64
	if code, payload, callErr := poller.agent.Do(ctx, server, http.MethodGet, "/v1/network/interfaces", nil); callErr == nil && code == 200 {
		var interfaces []struct {
			RX int64 `json:"rx_bytes"`
			TX int64 `json:"tx_bytes"`
		}
		if json.Unmarshal(payload, &interfaces) == nil {
			for _, item := range interfaces {
				rx += item.RX
				tx += item.TX
			}
		}
	}
	connections := 0
	if code, payload, callErr := poller.agent.Do(ctx, server, http.MethodGet, "/v1/network/connections", nil); callErr == nil && code == 200 {
		var items []json.RawMessage
		if json.Unmarshal(payload, &items) == nil {
			connections = len(items)
		}
	}
	txDB, err := poller.store.pool.Begin(ctx)
	if err != nil {
		return
	}
	defer txDB.Rollback(ctx)
	_, err = txDB.Exec(ctx, "UPDATE servers SET status='online',hostname=$2,last_seen=now(),updated_at=now() WHERE id=$1", server.ID, status.Hostname)
	if err == nil {
		_, err = txDB.Exec(ctx, `INSERT INTO metric_samples(server_id,cpu,memory,disk,rx,tx,connections,bans)
			VALUES($1,$2,$3,$4,$5,$6,$7,0)`, server.ID, status.CPU, memory, disk, rx, tx, connections)
	}
	if err == nil {
		err = txDB.Commit(ctx)
	}
	if err != nil {
		log.Printf("store metric for %s: %v", server.ID, err)
	}
	poller.pollLogs(ctx, server)
}

func (poller *Poller) pollLogs(ctx context.Context, server Server) {
	var cursor string
	var since time.Time
	err := poller.store.pool.QueryRow(ctx, `SELECT
		COALESCE((SELECT journal_cursor FROM agent_log_entries WHERE server_id=$1 ORDER BY event_at DESC,ingested_at DESC LIMIT 1),''),
		COALESCE((SELECT max(event_at) FROM agent_log_entries WHERE server_id=$1),now() - interval '2 minutes')`, server.ID).
		Scan(&cursor, &since)
	if err != nil {
		log.Printf("log cursor for %s: %v", server.ID, err)
		return
	}
	fallbackUsed := false
	for page := 0; page < 5; page++ {
		path := fmt.Sprintf("/v1/logs?since=%d&limit=200", since.Add(-time.Second).Unix())
		if cursor != "" {
			path = "/v1/logs?cursor=" + url.QueryEscape(cursor) + "&limit=200"
		}
		status, payload, callErr := poller.agent.Do(ctx, server, http.MethodGet, path, nil)
		if callErr != nil || status == http.StatusForbidden || status == http.StatusNotFound {
			return
		}
		if status != http.StatusOK {
			if cursor != "" && !fallbackUsed {
				cursor = ""
				fallbackUsed = true
				page--
				continue
			}
			log.Printf("log poll for %s returned %d", server.ID, status)
			return
		}
		var response struct {
			Entries []struct {
				Cursor    string  `json:"cursor"`
				Timestamp uint64  `json:"timestamp_unix_micros"`
				Unit      string  `json:"unit"`
				Priority  int     `json:"priority"`
				Message   string  `json:"message"`
				BootID    *string `json:"boot_id"`
				PID       *int    `json:"pid"`
			} `json:"entries"`
			Truncated bool `json:"truncated"`
		}
		if err = json.Unmarshal(payload, &response); err != nil {
			log.Printf("decode logs for %s: %v", server.ID, err)
			return
		}
		if len(response.Entries) > 200 {
			log.Printf("agent %s exceeded log entry bound", server.ID)
			return
		}
		tx, beginErr := poller.store.pool.Begin(ctx)
		if beginErr != nil {
			return
		}
		now := time.Now()
		for _, entry := range response.Entries {
			if len(entry.Cursor) == 0 || len(entry.Cursor) > 1024 || len(entry.Unit) > 128 ||
				len(entry.Message) > 65_536 || entry.Priority < 0 || entry.Priority > 7 {
				continue
			}
			eventAt := time.UnixMicro(int64(entry.Timestamp))
			if eventAt.After(now.Add(5*time.Minute)) || eventAt.Before(now.Add(-31*24*time.Hour)) {
				continue
			}
			_, err = tx.Exec(ctx, `INSERT INTO agent_log_entries
				(server_id,journal_cursor,event_at,unit,priority,message,boot_id,pid)
				VALUES($1,$2,$3,$4,$5,$6,$7,$8) ON CONFLICT DO NOTHING`, server.ID, entry.Cursor,
				eventAt, entry.Unit, entry.Priority, truncateUTF8(redactLogMessage(entry.Message), 65_536), entry.BootID, entry.PID)
			if err != nil {
				_ = tx.Rollback(ctx)
				return
			}
		}
		if err = tx.Commit(ctx); err != nil {
			log.Printf("store logs for %s: %v", server.ID, err)
			return
		}
		if len(response.Entries) == 0 || !response.Truncated {
			return
		}
		cursor = response.Entries[len(response.Entries)-1].Cursor
	}
	log.Printf("agent log backlog for %s exceeded 1000 entries in one poll", server.ID)
}

func redactLogMessage(message string) string {
	message = bearerSecret.ReplaceAllString(message, "${1}[REDACTED]")
	return namedSecret.ReplaceAllString(message, "${1}[REDACTED]")
}

func truncateUTF8(value string, maximum int) string {
	if len(value) <= maximum {
		return value
	}
	for maximum > 0 && value[maximum]&0xc0 == 0x80 {
		maximum--
	}
	return value[:maximum]
}

func percent(used, total uint64) float64 {
	if total == 0 {
		return 0
	}
	return float64(used) * 100 / float64(total)
}