package main

import (
	"bytes"
	"context"
	"crypto/sha256"
	"crypto/x509"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"slices"
	"strings"
	"sync"
	"time"

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

const ownershipLockKey int64 = 31987176956605541

const certificatePath = "/run/tls/tls.crt"
const equivalentCertificatePath = "/run/tls/./tls.crt"
const privateKeyPath = "/run/tls/tls.key"

var ownershipSessionSettings = []string{
	"set application_name = 'stalwart-controller'",
	"set tcp_keepalives_idle = 15",
	"set tcp_keepalives_interval = 5",
	"set tcp_keepalives_count = 3",
}

var errOwnerExists = errors.New("another controller owns reconciliation")
var errLeaderUnhealthy = errors.New("leader failed reconciliation repeatedly")
var buildRevision = "unknown"

type controllerStatus struct {
	Leader              bool   `json:"leader"`
	Ready               bool   `json:"ready"`
	PersistentFailure   bool   `json:"persistentFailure"`
	ConsecutiveFailures int    `json:"consecutiveFailures"`
	PlanRevision        string `json:"planRevision"`
	BuildRevision       string `json:"buildRevision"`
	LastAttempt         string `json:"lastAttempt,omitempty"`
	LastSuccess         string `json:"lastSuccess,omitempty"`
	NextAttempt         string `json:"nextAttempt,omitempty"`
	LastResult          string `json:"lastResult"`
}

type statusStore struct {
	mu     sync.RWMutex
	status controllerStatus
}

type ownershipConnection interface {
	Close(context.Context) error
	Ping(context.Context) error
}

type statusServer interface {
	Shutdown(context.Context) error
}

type controllerRuntime struct {
	acquireOwnership func(context.Context, string) (ownershipConnection, error)
	planRevision     func(string) (string, error)
	reconcile        func(context.Context, options) error
}

type certificateInventoryResult struct {
	ids     []string
	matches []string
}

func productionControllerRuntime() controllerRuntime {
	return controllerRuntime{
		acquireOwnership: func(ctx context.Context, dsnFile string) (ownershipConnection, error) {
			return acquireOwnership(ctx, dsnFile)
		},
		planRevision: planRevision,
		reconcile:    reconcileOnce,
	}
}

func closeOwnership(ownership ownershipConnection) {
	ownershipContext, cancelOwnership := context.WithTimeout(context.Background(), 5*time.Second)
	_ = ownership.Close(ownershipContext)
	cancelOwnership()
}

func shutdownStatusServer(server statusServer) {
	serverContext, cancelServer := context.WithTimeout(context.Background(), 5*time.Second)
	_ = server.Shutdown(serverContext)
	cancelServer()
}

func (s *statusStore) update(change func(*controllerStatus)) {
	s.mu.Lock()
	defer s.mu.Unlock()
	change(&s.status)
}

func (s *statusStore) snapshot() controllerStatus {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return s.status
}

func planRevision(path string) (string, error) {
	content, err := os.ReadFile(path)
	if err != nil {
		return "", err
	}
	sum := sha256.Sum256(content)
	return hex.EncodeToString(sum[:]), nil
}

func nextBackoff(current, maximum time.Duration) time.Duration {
	if current >= maximum/2 {
		return maximum
	}
	return current * 2
}

func metricBoolean(value bool) int {
	if value {
		return 1
	}
	return 0
}

func recordLeaderFailure(store *statusStore, threshold int) bool {
	relinquish := false
	store.update(func(status *controllerStatus) {
		status.ConsecutiveFailures++
		status.PersistentFailure = status.ConsecutiveFailures >= threshold
		status.Ready = status.LastSuccess != "" && !status.PersistentFailure
		status.LastResult = "failure"
		relinquish = status.PersistentFailure
	})
	return relinquish
}

func configureOwnershipSession(ctx context.Context, execute func(context.Context, string) error) error {
	for _, statement := range ownershipSessionSettings {
		if err := execute(ctx, statement); err != nil {
			return err
		}
	}
	return nil
}

func acquireOwnership(ctx context.Context, dsnFile string) (*pgx.Conn, error) {
	dsn, err := os.ReadFile(dsnFile)
	if err != nil {
		return nil, err
	}
	conn, err := pgx.Connect(ctx, strings.TrimSpace(string(dsn)))
	if err != nil {
		return nil, err
	}
	if err := configureOwnershipSession(ctx, func(ctx context.Context, statement string) error {
		_, err := conn.Exec(ctx, statement)
		return err
	}); err != nil {
		_ = conn.Close(ctx)
		return nil, err
	}
	var acquired bool
	err = conn.QueryRow(ctx, "select pg_try_advisory_lock($1)", ownershipLockKey).Scan(&acquired)
	if err != nil {
		_ = conn.Close(ctx)
		return nil, err
	}
	if !acquired {
		_ = conn.Close(ctx)
		return nil, errOwnerExists
	}
	return conn, nil
}

func secretPath(root, name string) (string, error) {
	if name == "" || filepath.IsAbs(name) || filepath.Clean(name) != name {
		return "", errors.New("invalid secret file name")
	}
	target := filepath.Join(root, name)
	relative, err := filepath.Rel(root, target)
	if err != nil || relative == ".." || strings.HasPrefix(relative, "../") {
		return "", errors.New("secret file escapes root")
	}
	return target, nil
}

func injectSecrets(value any, root string) (any, error) {
	switch typed := value.(type) {
	case []any:
		for index, item := range typed {
			replaced, err := injectSecrets(item, root)
			if err != nil {
				return nil, err
			}
			typed[index] = replaced
		}
		return typed, nil
	case map[string]any:
		if len(typed) == 1 {
			name, found := typed["$secretFile"].(string)
			asSet := false
			if !found {
				name, found = typed["$secretSet"].(string)
				asSet = found
			}
			if found {
				path, err := secretPath(root, name)
				if err != nil {
					return nil, err
				}
				content, err := os.ReadFile(path)
				if err != nil {
					return nil, err
				}
				resolved := strings.TrimSuffix(string(content), "\n")
				if asSet {
					return map[string]bool{resolved: true}, nil
				}
				return resolved, nil
			}
		}
		for key, item := range typed {
			replaced, err := injectSecrets(item, root)
			if err != nil {
				return nil, err
			}
			typed[key] = replaced
		}
		return typed, nil
	default:
		return value, nil
	}
}

func renderPlan(source, target, secretRoot string) error {
	content, err := os.ReadFile(source)
	if err != nil {
		return err
	}
	var operations []any
	if err := json.Unmarshal(content, &operations); err != nil {
		return err
	}
	replaced, err := injectSecrets(operations, secretRoot)
	if err != nil {
		return err
	}
	operations = replaced.([]any)
	file, err := os.OpenFile(target, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0o600)
	if err != nil {
		return err
	}
	if err := file.Chmod(0o600); err != nil {
		_ = file.Close()
		return err
	}
	encoder := json.NewEncoder(file)
	encoder.SetEscapeHTML(false)
	for _, operation := range operations {
		if err := encoder.Encode(operation); err != nil {
			_ = file.Close()
			return err
		}
	}
	return file.Close()
}

func runCLI(
	ctx context.Context,
	path string,
	environment []string,
	capture bool,
	arguments ...string,
) ([]byte, error) {
	command := exec.CommandContext(ctx, path, arguments...)
	if len(environment) != 0 {
		command.Env = mergedEnvironment(os.Environ(), environment)
	}
	var stdout bytes.Buffer
	if capture {
		command.Stdout = &stdout
	} else {
		command.Stdout = io.Discard
	}
	var stderr bytes.Buffer
	command.Stderr = &stderr
	if err := command.Run(); err != nil {
		action := strings.Join(arguments[:min(2, len(arguments))], " ")
		if ctx.Err() != nil {
			return nil, fmt.Errorf("stalwart CLI %s: %w", action, ctx.Err())
		}
		diagnostic := redactDiagnostic(stderr.String(), environment)
		if diagnostic == "" {
			return nil, fmt.Errorf("stalwart CLI %s: %w", action, err)
		}
		return nil, fmt.Errorf("stalwart CLI %s: %w: %s", action, err, diagnostic)
	}
	return stdout.Bytes(), nil
}

func redactDiagnostic(value string, environment []string) string {
	value = strings.TrimSpace(value)
	for _, item := range environment {
		_, secret, found := strings.Cut(item, "=")
		if found && secret != "" {
			value = strings.ReplaceAll(value, secret, "[redacted]")
		}
	}
	const maximum = 4096
	if len(value) > maximum {
		value = value[:maximum] + "..."
	}
	return value
}

func mergedEnvironment(base, overrides []string) []string {
	keys := make(map[string]struct{}, len(overrides))
	for _, item := range overrides {
		key, _, _ := strings.Cut(item, "=")
		keys[key] = struct{}{}
	}
	result := make([]string, 0, len(base)+len(overrides))
	for _, item := range base {
		key, _, _ := strings.Cut(item, "=")
		if _, replaced := keys[key]; !replaced {
			result = append(result, item)
		}
	}
	return append(result, overrides...)
}

func bootstrapEnvironment(value options) ([]string, error) {
	if value.bootstrapUser == "" {
		return nil, nil
	}
	password, err := os.ReadFile(value.bootstrapPasswordFile)
	if err != nil {
		return nil, err
	}
	return []string{
		"STALWART_USER=" + value.bootstrapUser,
		"STALWART_PASSWORD=" + strings.TrimSuffix(string(password), "\n"),
	}, nil
}

func credentialEnvironment(ctx context.Context, value options) ([]string, error) {
	environment, err := oidcEnvironment(ctx, value)
	if err != nil {
		return nil, fmt.Errorf("load OIDC credentials: %w", err)
	}
	if len(environment) == 0 {
		environment, err = bootstrapEnvironment(value)
		if err != nil {
			return nil, fmt.Errorf("load bootstrap credentials: %w", err)
		}
	}
	return append(environment, "STALWART_URL="+value.stalwartURL), nil
}

func certificateCoversHostname(subjectNames map[string]bool, hostname string) bool {
	certificate := x509.Certificate{}
	for name, enabled := range subjectNames {
		if enabled {
			certificate.DNSNames = append(certificate.DNSNames, name)
		}
	}
	return certificate.VerifyHostname(hostname) == nil
}

func certificateInventory(output []byte, san string) (certificateInventoryResult, error) {
	decoder := json.NewDecoder(bytes.NewReader(output))
	var result certificateInventoryResult
	for {
		var item struct {
			ID                      string          `json:"id"`
			SubjectAlternativeNames map[string]bool `json:"subjectAlternativeNames"`
		}
		err := decoder.Decode(&item)
		if errors.Is(err, io.EOF) {
			break
		}
		if err != nil {
			return certificateInventoryResult{}, err
		}
		if item.ID == "" {
			return certificateInventoryResult{}, errors.New("certificate query returned an empty id")
		}
		result.ids = append(result.ids, item.ID)
		if certificateCoversHostname(item.SubjectAlternativeNames, san) {
			result.matches = append(result.matches, item.ID)
		}
	}
	return result, nil
}

func queryCertificates(ctx context.Context, value options, environment []string) (certificateInventoryResult, error) {
	output, err := runCLI(
		ctx,
		value.cli,
		environment,
		true,
		"query",
		"Certificate",
		"--fields",
		"id,subjectAlternativeNames",
		"--json",
		"--no-color",
	)
	if err != nil {
		return certificateInventoryResult{}, err
	}
	return certificateInventory(output, value.certificateSAN)
}

func waitForCertificateSAN(ctx context.Context, value options, environment []string, id string) ([]string, error) {
	waitContext, cancel := context.WithTimeout(ctx, 30*time.Second)
	defer cancel()
	ticker := time.NewTicker(time.Second)
	defer ticker.Stop()
	for {
		inventory, err := queryCertificates(waitContext, value, environment)
		if err != nil {
			return nil, err
		}
		if slices.Contains(inventory.matches, id) {
			return inventory.ids, nil
		}
		select {
		case <-waitContext.Done():
			return nil, errors.New("updated certificate did not expose the declared SAN before timeout")
		case <-ticker.C:
		}
	}
}

func fileReferenceField(name, path string) string {
	return fmt.Sprintf(`%s={"@type":"File","filePath":%q}`, name, path)
}

func updateCertificate(ctx context.Context, value options, environment []string, id, path string) error {
	_, err := runCLI(
		ctx,
		value.cli,
		environment,
		false,
		"update",
		"Certificate",
		id,
		"--field",
		fileReferenceField("certificate", path),
		"--field",
		fileReferenceField("privateKey", privateKeyPath),
		"--no-color",
	)
	return err
}

func reconcileCertificate(ctx context.Context, value options, environment []string) error {
	inventory, err := queryCertificates(ctx, value, environment)
	if err != nil {
		return err
	}
	var id string
	if len(inventory.matches) == 0 {
		existing := make(map[string]struct{}, len(inventory.ids))
		for _, candidate := range inventory.ids {
			existing[candidate] = struct{}{}
		}
		if _, err := runCLI(
			ctx,
			value.cli,
			environment,
			false,
			"create",
			"Certificate",
			"--field",
			fileReferenceField("certificate", certificatePath),
			"--field",
			fileReferenceField("privateKey", privateKeyPath),
			"--no-color",
		); err != nil {
			return err
		}
		inventory, err = queryCertificates(ctx, value, environment)
		if err != nil {
			return err
		}
		if len(inventory.matches) == 0 {
			for _, candidate := range inventory.ids {
				if _, found := existing[candidate]; !found {
					inventory.matches = append(inventory.matches, candidate)
				}
			}
			if len(inventory.matches) != 1 {
				return errors.New("created certificate is not uniquely identifiable")
			}
		}
	}
	slices.Sort(inventory.matches)
	id = inventory.matches[0]
	// Stalwart derives metadata only when the certificate reference changes.
	// Round-tripping equivalent paths exposes rotations at a projected path.
	for _, path := range []string{equivalentCertificatePath, certificatePath} {
		if err := updateCertificate(ctx, value, environment, id, path); err != nil {
			return err
		}
	}
	if _, err := runCLI(ctx, value.cli, environment, false, "create", "Action/ReloadTlsCertificates", "--json", "{}", "--no-color"); err != nil {
		return err
	}
	ids, err := waitForCertificateSAN(ctx, value, environment, id)
	if err != nil {
		return err
	}
	if _, err := runCLI(
		ctx,
		value.cli,
		environment,
		false,
		"update",
		"SystemSettings",
		"--field",
		"defaultCertificateId="+id,
		"--no-color",
	); err != nil {
		return err
	}
	var obsolete []string
	for _, candidate := range ids {
		if candidate != id {
			obsolete = append(obsolete, candidate)
		}
	}
	slices.Sort(obsolete)
	if len(obsolete) != 0 {
		if _, err := runCLI(ctx, value.cli, environment, false, "delete", "Certificate", "--ids", strings.Join(obsolete, ","), "--no-color"); err != nil {
			return err
		}
		_, err = runCLI(ctx, value.cli, environment, false, "create", "Action/ReloadTlsCertificates", "--json", "{}", "--no-color")
		return err
	}
	return nil
}

func reconcileOnce(ctx context.Context, value options) error {
	if err := renderPlan(value.plan, value.renderedPlan, value.secretRoot); err != nil {
		return fmt.Errorf("render plan: %w", err)
	}
	defer func() {
		_ = os.Remove(value.renderedPlan)
	}()
	authenticationEnvironment, err := credentialEnvironment(ctx, value)
	if err != nil {
		return err
	}
	if _, err := runCLI(
		ctx,
		value.cli,
		authenticationEnvironment,
		false,
		"apply",
		"--file",
		value.renderedPlan,
		"--no-color",
	); err != nil {
		return err
	}
	if _, err := runCLI(
		ctx,
		value.cli,
		authenticationEnvironment,
		false,
		"create",
		"Action/ReloadSettings",
		"--json",
		"{}",
		"--no-color",
	); err != nil {
		return err
	}
	if value.updateApps {
		if _, err := runCLI(
			ctx,
			value.cli,
			authenticationEnvironment,
			false,
			"create",
			"Action/UpdateApps",
			"--json",
			"{}",
			"--no-color",
		); err != nil {
			return err
		}
	}
	if err := reconcileAppPasswords(ctx, value); err != nil {
		return err
	}
	if value.certificateSAN == "" {
		return nil
	}
	return reconcileCertificate(ctx, value, authenticationEnvironment)
}

func statusHandler(store *statusStore) http.Handler {
	mux := http.NewServeMux()
	mux.HandleFunc("/status", func(response http.ResponseWriter, _ *http.Request) {
		response.Header().Set("Content-Type", "application/json")
		_ = json.NewEncoder(response).Encode(store.snapshot())
	})
	mux.HandleFunc("/readyz", func(response http.ResponseWriter, _ *http.Request) {
		status := store.snapshot()
		if !status.Ready || status.PersistentFailure {
			response.WriteHeader(http.StatusServiceUnavailable)
			return
		}
		response.WriteHeader(http.StatusOK)
	})
	mux.HandleFunc("/healthz", func(response http.ResponseWriter, _ *http.Request) {
		response.WriteHeader(http.StatusOK)
	})
	mux.HandleFunc("/metrics", func(response http.ResponseWriter, _ *http.Request) {
		status := store.snapshot()
		response.Header().Set("Content-Type", "text/plain; version=0.0.4")
		_, _ = fmt.Fprintf(
			response,
			"stalwart_controller_leader %d\n"+
				"stalwart_controller_ready %d\n"+
				"stalwart_controller_persistent_failure %d\n"+
				"stalwart_controller_consecutive_failures %d\n",
			metricBoolean(status.Leader),
			metricBoolean(status.Ready),
			metricBoolean(status.PersistentFailure),
			status.ConsecutiveFailures,
		)
	})
	return mux
}

func waitForRetry(ctx context.Context, serverErrors <-chan error, delay time.Duration) error {
	timer := time.NewTimer(delay)
	defer timer.Stop()
	select {
	case <-ctx.Done():
		return ctx.Err()
	case err := <-serverErrors:
		return err
	case <-timer.C:
		return nil
	}
}

func runLeader(
	ctx context.Context,
	value options,
	store *statusStore,
	serverErrors <-chan error,
	ownership ownershipConnection,
	runtime controllerRuntime,
) error {
	defer closeOwnership(ownership)
	backoff := value.initialBackoff
	for {
		select {
		case err := <-serverErrors:
			return err
		default:
		}

		if err := ownership.Ping(ctx); err != nil {
			store.update(func(status *controllerStatus) {
				status.Leader = false
				status.Ready = false
				status.PersistentFailure = false
				status.ConsecutiveFailures = 0
				status.LastResult = "ownership-lost"
			})
			return errOwnerExists
		}
		revision, err := runtime.planRevision(value.plan)
		if err != nil {
			return err
		}
		attemptedAt := time.Now().UTC()
		store.update(func(status *controllerStatus) {
			status.PlanRevision = revision
			status.LastAttempt = attemptedAt.Format(time.RFC3339Nano)
			status.NextAttempt = ""
		})
		reconcileContext, cancel := context.WithTimeout(ctx, value.reconcileTimeout)
		err = runtime.reconcile(reconcileContext, value)
		cancel()
		if ctx.Err() != nil {
			return nil
		}

		var delay time.Duration
		if err == nil {
			succeededAt := time.Now().UTC()
			store.update(func(status *controllerStatus) {
				status.Ready = true
				status.PersistentFailure = false
				status.ConsecutiveFailures = 0
				status.LastSuccess = succeededAt.Format(time.RFC3339Nano)
				status.LastResult = "success"
			})
			backoff = value.initialBackoff
			delay = value.interval
		} else {
			fmt.Fprintln(os.Stderr, "reconciliation failed:", err)
			if recordLeaderFailure(store, value.persistentFailureAfter) {
				return errLeaderUnhealthy
			}
			delay = backoff
			backoff = nextBackoff(backoff, value.maxBackoff)
		}

		next := time.Now().UTC().Add(delay)
		store.update(func(status *controllerStatus) {
			status.NextAttempt = next.Format(time.RFC3339Nano)
		})
		if err := waitForRetry(ctx, serverErrors, delay); err != nil {
			return err
		}
	}
}

func runController(ctx context.Context, value options) error {
	return runControllerWithRuntime(ctx, value, productionControllerRuntime())
}

func runControllerWithRuntime(ctx context.Context, value options, runtime controllerRuntime) error {
	if err := validateOptions(value); err != nil {
		return err
	}
	if value.once {
		reconcileContext, cancel := context.WithTimeout(ctx, value.reconcileTimeout)
		defer cancel()
		return runtime.reconcile(reconcileContext, value)
	}
	revision, err := runtime.planRevision(value.plan)
	if err != nil {
		return err
	}
	store := &statusStore{
		status: controllerStatus{
			PlanRevision:  revision,
			BuildRevision: buildRevision,
			LastResult:    "standby",
		},
	}
	server := &http.Server{
		Addr:              value.statusAddress,
		Handler:           statusHandler(store),
		ReadHeaderTimeout: 5 * time.Second,
	}
	defer shutdownStatusServer(server)
	serverErrors := make(chan error, 1)
	go func() {
		errorValue := server.ListenAndServe()
		if !errors.Is(errorValue, http.ErrServerClosed) {
			serverErrors <- errorValue
		}
	}()
	backoff := value.initialBackoff
	for {
		standbyReady := false
		coolingDown := false
		ownershipContext, cancel := context.WithTimeout(ctx, value.reconcileTimeout)
		ownership, err := runtime.acquireOwnership(ownershipContext, value.dsnFile)
		cancel()
		if ctx.Err() != nil {
			return nil
		}
		if err == nil {
			store.update(func(status *controllerStatus) {
				status.Leader = true
				status.Ready = false
				status.PersistentFailure = false
				status.ConsecutiveFailures = 0
				status.LastResult = "not-attempted"
				status.NextAttempt = ""
			})
			err = runLeader(ctx, value, store, serverErrors, ownership, runtime)
			if ctx.Err() != nil {
				return nil
			}
			if errors.Is(err, errLeaderUnhealthy) {
				backoff = value.maxBackoff
				coolingDown = true
			} else if err != nil && !errors.Is(err, errOwnerExists) {
				return err
			} else {
				backoff = value.initialBackoff
			}
		} else if !errors.Is(err, errOwnerExists) {
			fmt.Fprintln(os.Stderr, "ownership acquisition failed:", err)
			backoff = nextBackoff(backoff, value.maxBackoff)
		} else {
			backoff = value.initialBackoff
			standbyReady = value.standbyReady
		}
		delay := backoff
		next := time.Now().UTC().Add(delay)
		store.update(func(status *controllerStatus) {
			status.Leader = false
			if coolingDown {
				status.Ready = false
				status.PersistentFailure = true
				status.LastResult = "cooldown"
			} else {
				status.Ready = standbyReady
				status.PersistentFailure = false
				status.ConsecutiveFailures = 0
				status.LastResult = "standby"
			}
			status.NextAttempt = next.Format(time.RFC3339Nano)
		})
		if err := waitForRetry(ctx, serverErrors, delay); err != nil {
			if errors.Is(err, context.Canceled) {
				return nil
			}
			return err
		}
	}
}