package actuator

import (
	"bytes"
	"context"
	"fmt"
	"log/slog"
	"net"
	"sync"
	"time"

	"golang.org/x/crypto/ssh"
)

var (
	hostKeyMu    sync.RWMutex
	hostKeyCache map[string]ssh.PublicKey
	hostKeyOnce sync.Once
	hostKeySrc   HostKeySource
)

// HostKeySource provides storage for SSH host public keys.
type HostKeySource interface {
	GetHostKey(ctx context.Context, hostname string) (string, error)
	SetHostKey(ctx context.Context, hostname string, key string) error
}

// SetHostKeySource sets the host key source. Must be called before
// any SSH connections. A nil source enables TOFU-only mode (keys
// accepted in memory but not persisted).
func SetHostKeySource(src HostKeySource) {
	hostKeyMu.Lock()
	defer hostKeyMu.Unlock()
	hostKeySrc = src
}

// HostKeyCallback returns an ssh.HostKeyCallback that verifies host keys.
// Known keys are verified (MITM detection). Unknown keys are accepted
// via TOFU and optionally persisted to the source.
func HostKeyCallback() ssh.HostKeyCallback {
	return hostKeyVerify
}

func hostKeyVerify(hostname string, remote net.Addr, key ssh.PublicKey) error {
	hostKeyOnce.Do(func() {
		hostKeyCache = make(map[string]ssh.PublicKey)
	})

	normalized := hostWithoutPort(hostname)

	hostKeyMu.RLock()
	known, exists := hostKeyCache[normalized]
	hostKeyMu.RUnlock()

	if exists {
		if bytes.Equal(key.Marshal(), known.Marshal()) {
			return nil
		}
		return fmt.Errorf("SSH HOST KEY CHANGED for %s (possible MITM)", normalized)
	}

	hostKeyMu.Lock()
	hostKeyCache[normalized] = key
	hostKeyMu.Unlock()

	slog.Info("ssh: accepting new host key (TOFU)", "host", normalized)

	if hostKeySrc != nil {
		go persistHostKey(normalized, key)
	}

	return nil
}

func persistHostKey(hostname string, key ssh.PublicKey) {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	keyBase64 := key.Type() + " " + string(key.Marshal())
	if err := hostKeySrc.SetHostKey(ctx, "ssh/host-keys/"+hostname, keyBase64); err != nil {
		slog.Warn("ssh: failed to persist host key", "host", hostname, "error", err)
	}
}

// LoadHostKeys pre-loads known host keys from the source into the
// in-memory cache. Call at startup to avoid TOFU on first connection.
// The source should return key lines in the format "key-type base64-data".
func LoadHostKeys(ctx context.Context, hostnames []string, src HostKeySource) {
	if src == nil {
		return
	}
	SetHostKeySource(src)

	hostKeyMu.Lock()
	defer hostKeyMu.Unlock()

	if hostKeyCache == nil {
		hostKeyCache = make(map[string]ssh.PublicKey)
	}

	loaded := 0
	for _, hostname := range hostnames {
		keyData, err := src.GetHostKey(ctx, "ssh/host-keys/"+hostname)
		if err != nil {
			slog.Debug("ssh: no stored key for host", "host", hostname, "error", err)
			continue
		}
		pubKey, _, _, _, err := ssh.ParseAuthorizedKey([]byte(keyData))
		if err != nil {
			slog.Warn("ssh: invalid stored key for host", "host", hostname, "error", err)
			continue
		}
		hostKeyCache[hostname] = pubKey
		loaded++
	}

	if loaded > 0 {
		slog.Info("ssh: loaded host keys from Infisical", "count", loaded)
	}
}

func hostWithoutPort(hostname string) string {
	for i := len(hostname) - 1; i >= 0; i-- {
		if hostname[i] == ':' {
			return hostname[:i]
		}
	}
	return hostname
}