package sshx

import (
	"context"
	"errors"
	"fmt"
	"net"
	"os"
	"path/filepath"
	"sync"
	"time"

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

// KnownHosts is the trust store for target host keys. It behaves like OpenSSH:
// an unknown key is refused until the operator confirms the fingerprint, and a
// changed key is refused outright.
type KnownHosts struct {
	path string
	mu   sync.Mutex
}

// NewKnownHosts opens (and creates if needed) the store at path.
func NewKnownHosts(path string) (*KnownHosts, error) {
	if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil {
		return nil, fmt.Errorf("create key store directory: %w", err)
	}
	f, err := os.OpenFile(path, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0o600)
	if err != nil {
		return nil, fmt.Errorf("open known hosts file: %w", err)
	}
	f.Close()
	return &KnownHosts{path: path}, nil
}

// Path returns the on-disk location of the store.
func (k *KnownHosts) Path() string { return k.path }

// Check implements the ssh.HostKeyCallback contract.
func (k *KnownHosts) Check(hostname string, remote net.Addr, key ssh.PublicKey) error {
	k.mu.Lock()
	defer k.mu.Unlock()

	cb, err := knownhosts.New(k.path)
	if err != nil {
		return fmt.Errorf("read known hosts: %w", err)
	}
	err = cb(hostname, remote, key)
	if err == nil {
		return nil
	}
	var keyErr *knownhosts.KeyError
	if errors.As(err, &keyErr) {
		return &HostKeyError{
			Host:        hostname,
			Fingerprint: ssh.FingerprintSHA256(key),
			KeyType:     key.Type(),
			Changed:     len(keyErr.Want) > 0,
		}
	}
	return err
}

// Trust records a host key so later connections succeed.
func (k *KnownHosts) Trust(hostname string, key ssh.PublicKey) error {
	k.mu.Lock()
	defer k.mu.Unlock()

	f, err := os.OpenFile(k.path, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0o600)
	if err != nil {
		return fmt.Errorf("open known hosts for write: %w", err)
	}
	defer f.Close()
	line := knownhosts.Line([]string{knownhosts.Normalize(hostname)}, key)
	if _, err := f.WriteString(line + "\n"); err != nil {
		return fmt.Errorf("record host key: %w", err)
	}
	return nil
}

// Forget removes every entry for a host, so a changed key can be re-approved.
func (k *KnownHosts) Forget(hostname string) error {
	k.mu.Lock()
	defer k.mu.Unlock()

	b, err := os.ReadFile(k.path)
	if err != nil {
		return err
	}
	want := knownhosts.Normalize(hostname)
	var kept []byte
	for _, line := range splitLines(b) {
		if len(line) == 0 || line[0] == '#' {
			kept = append(kept, line...)
			kept = append(kept, '\n')
			continue
		}
		_, hosts, _, _, _, perr := ssh.ParseKnownHosts(append(line, '\n'))
		if perr == nil && containsHost(hosts, want) {
			continue
		}
		kept = append(kept, line...)
		kept = append(kept, '\n')
	}
	return os.WriteFile(k.path, kept, 0o600)
}

// HostKeyInfo is the fingerprint presented by a host, shown to the operator
// before they decide to trust it.
type HostKeyInfo struct {
	Host        string `json:"host"`
	KeyType     string `json:"keyType"`
	Fingerprint string `json:"fingerprint"`
	Trusted     bool   `json:"trusted"`
	Changed     bool   `json:"changed"`
}

// Probe opens a TCP connection just far enough to read the host key, without
// authenticating. Used by the "check fingerprint" step in the UI.
func Probe(ctx context.Context, cfg Config, hk *KnownHosts) (*HostKeyInfo, error) {
	timeout := cfg.Timeout
	if timeout == 0 {
		timeout = 15 * time.Second
	}
	var captured ssh.PublicKey
	clientCfg := &ssh.ClientConfig{
		User:    cfg.User,
		Timeout: timeout,
		HostKeyCallback: func(_ string, _ net.Addr, key ssh.PublicKey) error {
			captured = key
			// Stop the handshake here: reading the key is all this needs.
			return errProbeDone
		},
	}
	d := net.Dialer{Timeout: timeout}
	conn, err := d.DialContext(ctx, "tcp", cfg.addr())
	if err != nil {
		return nil, fmt.Errorf("connect to %s: %w", cfg.addr(), err)
	}
	defer conn.Close()
	_, _, _, err = ssh.NewClientConn(conn, cfg.addr(), clientCfg)
	if captured == nil {
		return nil, fmt.Errorf("read host key from %s: %w", cfg.addr(), err)
	}

	info := &HostKeyInfo{
		Host:        cfg.addr(),
		KeyType:     captured.Type(),
		Fingerprint: ssh.FingerprintSHA256(captured),
	}
	switch checkErr := hk.Check(cfg.addr(), conn.RemoteAddr(), captured).(type) {
	case nil:
		info.Trusted = true
	case *HostKeyError:
		info.Changed = checkErr.Changed
	}
	return info, nil
}

// TrustFromProbe re-reads the host key and stores it. Taking the key from a
// fresh handshake rather than from client-supplied input means the UI can only
// approve a fingerprint it actually saw.
func TrustFromProbe(ctx context.Context, cfg Config, hk *KnownHosts, expectFingerprint string) error {
	info, err := Probe(ctx, cfg, hk)
	if err != nil {
		return err
	}
	if expectFingerprint != "" && info.Fingerprint != expectFingerprint {
		return fmt.Errorf("host key changed between check and approval (%s vs %s); aborting",
			expectFingerprint, info.Fingerprint)
	}
	var captured ssh.PublicKey
	clientCfg := &ssh.ClientConfig{
		User:    cfg.User,
		Timeout: 15 * time.Second,
		HostKeyCallback: func(_ string, _ net.Addr, key ssh.PublicKey) error {
			captured = key
			return errProbeDone
		},
	}
	conn, err := net.DialTimeout("tcp", cfg.addr(), 15*time.Second)
	if err != nil {
		return err
	}
	defer conn.Close()
	_, _, _, _ = ssh.NewClientConn(conn, cfg.addr(), clientCfg)
	if captured == nil {
		return errors.New("could not read host key")
	}
	if ssh.FingerprintSHA256(captured) != info.Fingerprint {
		return errors.New("host key is unstable; aborting")
	}
	if info.Changed {
		if err := hk.Forget(cfg.addr()); err != nil {
			return fmt.Errorf("drop previous host key: %w", err)
		}
	}
	return hk.Trust(cfg.addr(), captured)
}

var errProbeDone = errors.New("host key captured")

func splitLines(b []byte) [][]byte {
	var out [][]byte
	start := 0
	for i := 0; i < len(b); i++ {
		if b[i] == '\n' {
			line := b[start:i]
			if n := len(line); n > 0 && line[n-1] == '\r' {
				line = line[:n-1]
			}
			out = append(out, line)
			start = i + 1
		}
	}
	if start < len(b) {
		out = append(out, b[start:])
	}
	return out
}

func containsHost(hosts []string, want string) bool {
	for _, h := range hosts {
		if knownhosts.Normalize(h) == want {
			return true
		}
	}
	return false
}