package main

import (
	"context"
	"crypto/rand"
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"os"
	"sort"
	"strconv"
	"strings"
	"time"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
)

const schema = `
CREATE TABLE IF NOT EXISTS users (
  id text PRIMARY KEY,
  username text UNIQUE NOT NULL,
  password_hash text NOT NULL,
  role text NOT NULL CHECK (role IN ('viewer','operator','administrator')),
  disabled boolean NOT NULL DEFAULT false,
  created_at timestamptz NOT NULL DEFAULT now(),
  updated_at timestamptz NOT NULL DEFAULT now()
);
CREATE TABLE IF NOT EXISTS sessions (
  id_hash bytea PRIMARY KEY,
  user_id text NOT NULL REFERENCES users(id) ON DELETE CASCADE,
  csrf_hash bytea NOT NULL,
  source_ip text NOT NULL,
  expires_at timestamptz NOT NULL,
  created_at timestamptz NOT NULL DEFAULT now()
);
CREATE INDEX IF NOT EXISTS sessions_expires_idx ON sessions(expires_at);
CREATE TABLE IF NOT EXISTS servers (
  id text PRIMARY KEY,
  name text NOT NULL,
  hostname text NOT NULL DEFAULT '',
  agent_endpoint text NOT NULL,
  credential_ref text NOT NULL,
  tags jsonb NOT NULL DEFAULT '[]'::jsonb,
  status text NOT NULL DEFAULT 'unknown',
  last_seen timestamptz,
  created_at timestamptz NOT NULL DEFAULT now(),
  updated_at timestamptz NOT NULL DEFAULT now()
);
CREATE TABLE IF NOT EXISTS metric_samples (
  server_id text NOT NULL REFERENCES servers(id) ON DELETE CASCADE,
  sampled_at timestamptz NOT NULL DEFAULT now(),
  cpu double precision NOT NULL,
  memory double precision NOT NULL,
  disk double precision NOT NULL,
  rx bigint NOT NULL,
  tx bigint NOT NULL,
  connections integer NOT NULL,
  bans integer NOT NULL,
  PRIMARY KEY(server_id, sampled_at)
);
CREATE INDEX IF NOT EXISTS metric_samples_time_idx ON metric_samples(sampled_at);
CREATE TABLE IF NOT EXISTS agent_log_entries (
  server_id text NOT NULL REFERENCES servers(id) ON DELETE CASCADE,
  journal_cursor text NOT NULL,
  event_at timestamptz NOT NULL,
  unit text NOT NULL,
  priority smallint NOT NULL CHECK (priority BETWEEN 0 AND 7),
  message text NOT NULL,
  boot_id text,
  pid integer,
  ingested_at timestamptz NOT NULL DEFAULT now(),
  PRIMARY KEY(server_id, journal_cursor)
);
CREATE INDEX IF NOT EXISTS agent_log_entries_time_idx ON agent_log_entries(event_at DESC);
CREATE INDEX IF NOT EXISTS agent_log_entries_server_time_idx ON agent_log_entries(server_id, event_at DESC);
CREATE TABLE IF NOT EXISTS audit_events (
  id bigserial PRIMARY KEY,
  user_id text REFERENCES users(id) ON DELETE SET NULL,
  username text NOT NULL,
  server_id text REFERENCES servers(id) ON DELETE SET NULL,
  action text NOT NULL,
  target text NOT NULL,
  before_state jsonb,
  after_state jsonb,
  result text NOT NULL,
  source_ip text NOT NULL,
  created_at timestamptz NOT NULL DEFAULT now()
);
CREATE INDEX IF NOT EXISTS audit_events_time_idx ON audit_events(created_at DESC);
`

type Store struct{ pool *pgxpool.Pool }

type Server struct {
	ID            string          `json:"id"`
	Name          string          `json:"name"`
	Hostname      string          `json:"hostname"`
	AgentEndpoint string          `json:"agent_endpoint"`
	CredentialRef string          `json:"-"`
	Tags          json.RawMessage `json:"tags"`
	Status        string          `json:"status"`
	LastSeen      *time.Time      `json:"last_seen"`
	CreatedAt     time.Time       `json:"created_at"`
}

type AuditEvent struct {
	ID        int64           `json:"id"`
	Username  string          `json:"username"`
	ServerID  *string         `json:"server_id"`
	Action    string          `json:"action"`
	Target    string          `json:"target"`
	Before    json.RawMessage `json:"before"`
	After     json.RawMessage `json:"after"`
	Result    string          `json:"result"`
	SourceIP  string          `json:"source_ip"`
	CreatedAt time.Time       `json:"created_at"`
}

type MetricSample struct {
	SampledAt   time.Time `json:"sampled_at"`
	CPU         float64   `json:"cpu"`
	Memory      float64   `json:"memory"`
	Disk        float64   `json:"disk"`
	RX          int64     `json:"rx"`
	TX          int64     `json:"tx"`
	Connections int       `json:"connections"`
	Bans        int       `json:"bans"`
}

type AgentLogEntry struct {
	ServerID  string    `json:"server_id"`
	Timestamp time.Time `json:"timestamp"`
	Unit      string    `json:"unit"`
	Priority  int       `json:"priority"`
	Message   string    `json:"message"`
	BootID    *string   `json:"boot_id,omitempty"`
	PID       *int      `json:"pid,omitempty"`
}

type MetricPoint struct {
	ServerID   string
	ServerName string
	Timestamp  time.Time
	Value      float64
}

type IntegrationLog struct {
	ServerID   string
	ServerName string
	Timestamp  time.Time
	Unit       string
	Priority   int
	Message    string
}

func OpenStore(ctx context.Context, databaseURL string) (*Store, error) {
	var lastErr error
	for attempt := 0; attempt < 30; attempt++ {
		pool, err := pgxpool.New(ctx, databaseURL)
		if err == nil {
			err = pool.Ping(ctx)
		}
		if err == nil {
			store := &Store{pool: pool}
			if _, err = pool.Exec(ctx, schema); err == nil {
				return store, nil
			}
			pool.Close()
		}
		lastErr = err
		select {
		case <-ctx.Done():
			return nil, ctx.Err()
		case <-time.After(time.Second):
		}
	}
	return nil, fmt.Errorf("connect/migrate database: %w", lastErr)
}

func (store *Store) Close() { store.pool.Close() }

func (store *Store) Bootstrap(ctx context.Context, config Config) error {
	var count int
	if err := store.pool.QueryRow(ctx, "SELECT count(*) FROM users").Scan(&count); err != nil {
		return err
	}
	if count == 0 {
		if !validUsername(config.AdminUsername) || config.AdminPasswordFile == "" {
			return errors.New("first start requires INITIAL_ADMIN_USERNAME and INITIAL_ADMIN_PASSWORD_FILE")
		}
		password, err := os.ReadFile(config.AdminPasswordFile)
		if err != nil {
			return fmt.Errorf("read initial admin password: %w", err)
		}
		hash, err := HashPassword(strings.TrimSpace(string(password)))
		if err != nil {
			return fmt.Errorf("initial admin password: %w", err)
		}
		if _, err = store.pool.Exec(ctx,
			"INSERT INTO users(id, username, password_hash, role) VALUES($1,$2,$3,$4)",
			newID(), config.AdminUsername, hash, RoleAdministrator); err != nil {
			return err
		}
		log.Printf("created initial administrator %q", config.AdminUsername)
	}
	if config.BootstrapEndpoint != "" {
		_, err := store.pool.Exec(ctx, `INSERT INTO servers(id,name,agent_endpoint,credential_ref)
			SELECT $1,$2,$3,$4 WHERE NOT EXISTS (SELECT 1 FROM servers)`,
			newID(), config.BootstrapServerName, config.BootstrapEndpoint, config.BootstrapCredential)
		return err
	}
	return nil
}

func (store *Store) LoginUser(ctx context.Context, username string) (User, string, bool, error) {
	var user User
	var hash string
	var disabled bool
	err := store.pool.QueryRow(ctx,
		"SELECT id,username,role,password_hash,disabled FROM users WHERE username=$1", username).
		Scan(&user.ID, &user.Username, &user.Role, &hash, &disabled)
	if errors.Is(err, pgx.ErrNoRows) {
		return User{}, "", false, nil
	}
	return user, hash, disabled, err
}

func (store *Store) CreateSession(ctx context.Context, userID, sourceIP string, ttl time.Duration) (string, string, error) {
	session, sessionHash, err := randomToken()
	if err != nil {
		return "", "", err
	}
	csrf, csrfHash, err := randomToken()
	if err != nil {
		return "", "", err
	}
	_, err = store.pool.Exec(ctx,
		"INSERT INTO sessions(id_hash,user_id,csrf_hash,source_ip,expires_at) VALUES($1,$2,$3,$4,$5)",
		sessionHash, userID, csrfHash, sourceIP, time.Now().Add(ttl))
	return session, csrf, err
}

func (store *Store) SessionUser(ctx context.Context, token string) (User, []byte, error) {
	var user User
	var csrfHash []byte
	err := store.pool.QueryRow(ctx, `SELECT u.id,u.username,u.role,s.csrf_hash
		FROM sessions s JOIN users u ON u.id=s.user_id
		WHERE s.id_hash=$1 AND s.expires_at>now() AND NOT u.disabled`, tokenHash(token)).
		Scan(&user.ID, &user.Username, &user.Role, &csrfHash)
	return user, csrfHash, err
}

func (store *Store) DeleteSession(ctx context.Context, token string) error {
	_, err := store.pool.Exec(ctx, "DELETE FROM sessions WHERE id_hash=$1", tokenHash(token))
	return err
}

func (store *Store) Servers(ctx context.Context) ([]Server, error) {
	rows, err := store.pool.Query(ctx, `SELECT id,name,hostname,agent_endpoint,credential_ref,tags,status,last_seen,created_at
		FROM servers ORDER BY name`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	servers := []Server{}
	for rows.Next() {
		var server Server
		if err = rows.Scan(&server.ID, &server.Name, &server.Hostname, &server.AgentEndpoint,
			&server.CredentialRef, &server.Tags, &server.Status, &server.LastSeen, &server.CreatedAt); err != nil {
			return nil, err
		}
		servers = append(servers, server)
	}
	return servers, rows.Err()
}

func (store *Store) Server(ctx context.Context, id string) (Server, error) {
	var server Server
	err := store.pool.QueryRow(ctx, `SELECT id,name,hostname,agent_endpoint,credential_ref,tags,status,last_seen,created_at
		FROM servers WHERE id=$1`, id).Scan(&server.ID, &server.Name, &server.Hostname, &server.AgentEndpoint,
		&server.CredentialRef, &server.Tags, &server.Status, &server.LastSeen, &server.CreatedAt)
	return server, err
}

func (store *Store) CreateServer(ctx context.Context, server Server) error {
	_, err := store.pool.Exec(ctx, `INSERT INTO servers(id,name,hostname,agent_endpoint,credential_ref,tags)
		VALUES($1,$2,$3,$4,$5,$6)`, server.ID, server.Name, server.Hostname,
		server.AgentEndpoint, server.CredentialRef, server.Tags)
	return err
}

func (store *Store) Audit(ctx context.Context, user *User, serverID, action, target string,
	before, after any, result, sourceIP string) {
	var userID any
	username := "anonymous"
	if user != nil {
		userID, username = user.ID, user.Username
	}
	beforeJSON, _ := json.Marshal(before)
	afterJSON, _ := json.Marshal(after)
	if before == nil {
		beforeJSON = nil
	}
	if after == nil {
		afterJSON = nil
	}
	if _, err := store.pool.Exec(ctx, `INSERT INTO audit_events
		(user_id,username,server_id,action,target,before_state,after_state,result,source_ip)
		VALUES($1,$2,NULLIF($3,''),$4,$5,$6,$7,$8,$9)`, userID, username, serverID,
		action, target, beforeJSON, afterJSON, result, sourceIP); err != nil {
		log.Printf("write audit event: %v", err)
	}
}

func (store *Store) AuditEvents(ctx context.Context, limit int) ([]AuditEvent, error) {
	rows, err := store.pool.Query(ctx, `SELECT id,username,server_id,action,target,
		COALESCE(before_state,'null'::jsonb),COALESCE(after_state,'null'::jsonb),result,source_ip,created_at
		FROM audit_events ORDER BY created_at DESC LIMIT $1`, limit)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	events := []AuditEvent{}
	for rows.Next() {
		var event AuditEvent
		if err = rows.Scan(&event.ID, &event.Username, &event.ServerID, &event.Action, &event.Target,
			&event.Before, &event.After, &event.Result, &event.SourceIP, &event.CreatedAt); err != nil {
			return nil, err
		}
		events = append(events, event)
	}
	return events, rows.Err()
}

func (store *Store) Metrics(ctx context.Context, serverID string, limit int) ([]MetricSample, error) {
	rows, err := store.pool.Query(ctx, `SELECT sampled_at,cpu,memory,disk,rx,tx,connections,bans
		FROM metric_samples WHERE server_id=$1 ORDER BY sampled_at DESC LIMIT $2`, serverID, limit)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	samples := []MetricSample{}
	for rows.Next() {
		var sample MetricSample
		if err = rows.Scan(&sample.SampledAt, &sample.CPU, &sample.Memory, &sample.Disk,
			&sample.RX, &sample.TX, &sample.Connections, &sample.Bans); err != nil {
			return nil, err
		}
		samples = append(samples, sample)
	}
	return samples, rows.Err()
}

func (store *Store) Logs(ctx context.Context, serverID string, limit int) ([]AgentLogEntry, error) {
	rows, err := store.pool.Query(ctx, `SELECT server_id,event_at,unit,priority,message,boot_id,pid
		FROM agent_log_entries WHERE server_id=$1 ORDER BY event_at DESC LIMIT $2`, serverID, limit)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	entries := []AgentLogEntry{}
	for rows.Next() {
		var entry AgentLogEntry
		if err = rows.Scan(&entry.ServerID, &entry.Timestamp, &entry.Unit, &entry.Priority,
			&entry.Message, &entry.BootID, &entry.PID); err != nil {
			return nil, err
		}
		entries = append(entries, entry)
	}
	return entries, rows.Err()
}

func (store *Store) MetricSeries(ctx context.Context) ([]map[string]string, error) {
	rows, err := store.pool.Query(ctx, "SELECT id,name FROM servers ORDER BY id LIMIT 1000")
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	servers := [][2]string{}
	for rows.Next() {
		var server [2]string
		if err = rows.Scan(&server[0], &server[1]); err != nil {
			return nil, err
		}
		servers = append(servers, server)
	}
	if err = rows.Err(); err != nil {
		return nil, err
	}
	names := make([]string, 0, len(metricColumns))
	for name := range metricColumns {
		names = append(names, name)
	}
	sort.Strings(names)
	series := make([]map[string]string, 0, len(servers)*len(names))
	for _, server := range servers {
		for _, name := range names {
			series = append(series, map[string]string{
				"__name__": name, "server_id": server[0], "server_name": server[1],
			})
		}
	}
	return series, nil
}

func (store *Store) MetricLabelValues(ctx context.Context, label string) ([]string, error) {
	column := ""
	switch label {
	case "server_id":
		column = "id"
	case "server_name":
		column = "name"
	default:
		return nil, errors.New("unsupported metric label")
	}
	rows, err := store.pool.Query(ctx, "SELECT "+column+" FROM servers ORDER BY "+column+" LIMIT 10000")
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	values := []string{}
	for rows.Next() {
		var value string
		if err = rows.Scan(&value); err != nil {
			return nil, err
		}
		values = append(values, value)
	}
	return values, rows.Err()
}

func (store *Store) InstantMetric(ctx context.Context, query metricQuery, at time.Time) ([]MetricPoint, error) {
	column, found := metricColumns[query.Name]
	if !found {
		return nil, errors.New("unsupported metric")
	}
	statement := fmt.Sprintf(`SELECT DISTINCT ON (m.server_id)
		m.server_id,s.name,m.sampled_at,m.%s::double precision
		FROM metric_samples m JOIN servers s ON s.id=m.server_id
		WHERE m.sampled_at <= $1 AND ($2='' OR s.id=$2) AND ($3='' OR s.name=$3)
		ORDER BY m.server_id,m.sampled_at DESC LIMIT 10000`, column)
	rows, err := store.pool.Query(ctx, statement, at, query.ServerID, query.ServerName)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	return scanMetricPoints(rows)
}

func (store *Store) RangeMetric(ctx context.Context, query metricQuery, start, end time.Time, step int) ([]MetricPoint, error) {
	column, found := metricColumns[query.Name]
	if !found {
		return nil, errors.New("unsupported metric")
	}
	statement := fmt.Sprintf(`SELECT m.server_id,s.name,
		to_timestamp(floor(extract(epoch FROM m.sampled_at)/$3)*$3) AS bucket,
		avg(m.%s::double precision)
		FROM metric_samples m JOIN servers s ON s.id=m.server_id
		WHERE m.sampled_at BETWEEN $1 AND $2 AND ($4='' OR s.id=$4) AND ($5='' OR s.name=$5)
		GROUP BY m.server_id,s.name,bucket ORDER BY m.server_id,bucket LIMIT 100000`, column)
	rows, err := store.pool.Query(ctx, statement, start, end, step, query.ServerID, query.ServerName)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	return scanMetricPoints(rows)
}

func scanMetricPoints(rows pgx.Rows) ([]MetricPoint, error) {
	points := []MetricPoint{}
	for rows.Next() {
		var point MetricPoint
		if err := rows.Scan(&point.ServerID, &point.ServerName, &point.Timestamp, &point.Value); err != nil {
			return nil, err
		}
		points = append(points, point)
	}
	return points, rows.Err()
}

func (store *Store) LogLabelValues(ctx context.Context, label string) ([]string, error) {
	column := ""
	switch label {
	case "server_id":
		column = "s.id"
	case "server_name":
		column = "s.name"
	case "unit":
		column = "l.unit"
	case "priority":
		column = "l.priority::text"
	default:
		return nil, errors.New("unsupported log label")
	}
	statement := `SELECT DISTINCT ` + column + ` AS value FROM agent_log_entries l
		JOIN servers s ON s.id=l.server_id ORDER BY value LIMIT 10000`
	rows, err := store.pool.Query(ctx, statement)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	values := []string{}
	for rows.Next() {
		var value string
		if err = rows.Scan(&value); err != nil {
			return nil, err
		}
		values = append(values, value)
	}
	return values, rows.Err()
}

func (store *Store) LogSeries(ctx context.Context) ([]map[string]string, error) {
	rows, err := store.pool.Query(ctx, `SELECT DISTINCT s.id,s.name,l.unit,l.priority
		FROM agent_log_entries l JOIN servers s ON s.id=l.server_id
		ORDER BY s.id,l.unit,l.priority LIMIT 10000`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	series := []map[string]string{}
	for rows.Next() {
		var serverID, serverName, unit string
		var priority int
		if err = rows.Scan(&serverID, &serverName, &unit, &priority); err != nil {
			return nil, err
		}
		series = append(series, map[string]string{
			"server_id": serverID, "server_name": serverName, "unit": unit,
			"priority": strconv.Itoa(priority),
		})
	}
	return series, rows.Err()
}

func (store *Store) IntegrationLogs(ctx context.Context, start, end time.Time, limit int, direction string) ([]IntegrationLog, error) {
	order := "DESC"
	if direction == "forward" {
		order = "ASC"
	}
	statement := `SELECT l.server_id,s.name,l.event_at,l.unit,l.priority,l.message
		FROM agent_log_entries l JOIN servers s ON s.id=l.server_id
		WHERE l.event_at BETWEEN $1 AND $2 ORDER BY l.event_at ` + order + ` LIMIT $3`
	rows, err := store.pool.Query(ctx, statement, start, end, limit)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	entries := []IntegrationLog{}
	for rows.Next() {
		var entry IntegrationLog
		if err = rows.Scan(&entry.ServerID, &entry.ServerName, &entry.Timestamp, &entry.Unit,
			&entry.Priority, &entry.Message); err != nil {
			return nil, err
		}
		entries = append(entries, entry)
	}
	return entries, rows.Err()
}

func newID() string {
	var value [16]byte
	if _, err := rand.Read(value[:]); err != nil {
		panic("system random source unavailable")
	}
	value[6] = (value[6] & 0x0f) | 0x40
	value[8] = (value[8] & 0x3f) | 0x80
	return fmt.Sprintf("%08x-%04x-%04x-%04x-%012x",
		value[0:4], value[4:6], value[6:8], value[8:10], value[10:16])
}