package wgsync

// Group C — the SSH push client against an IN-PROCESS x/crypto/ssh server (Scenarios A/B/C-c4).
// The server captures the exact stdin payload, so tests assert the pushed BYTES; the host-key
// tests prove the pin both ways (correct key accepted, wrong key refused).

import (
	"bytes"
	"context"
	"crypto/ecdsa"
	"crypto/ed25519"
	"crypto/elliptic"
	"crypto/rand"
	"encoding/pem"
	"io"
	"log"
	"net"
	"strings"
	"sync"
	"testing"
	"time"

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

// testKeys generates an ed25519 keypair and returns (PEM private key, ssh.Signer).
func testKeys(t *testing.T) ([]byte, ssh.Signer) {
	t.Helper()
	_, priv, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		t.Fatalf("ed25519: %v", err)
	}
	block, err := ssh.MarshalPrivateKey(priv, "")
	if err != nil {
		t.Fatalf("MarshalPrivateKey: %v", err)
	}
	signer, err := ssh.NewSignerFromKey(priv)
	if err != nil {
		t.Fatalf("NewSignerFromKey: %v", err)
	}
	return pem.EncodeToMemory(block), signer
}

// testServer is a single-shot in-process SSH server handling one exec session.
type testServer struct {
	addr     string
	mu       sync.Mutex
	captured []byte // stdin the "script" received
	cmd      string // the exec command string requested
}

// startTestServer runs an SSH server that accepts clientSigner's key, serves with hostSigner,
// reads all stdin, replies stdoutResp/stderrResp and exitStatus.
func startTestServer(t *testing.T, hostSigner ssh.Signer, clientSigner ssh.Signer,
	stdoutResp, stderrResp string, exitStatus uint32) *testServer {
	t.Helper()
	cfg := &ssh.ServerConfig{
		PublicKeyCallback: func(conn ssh.ConnMetadata, key ssh.PublicKey) (*ssh.Permissions, error) {
			if bytes.Equal(key.Marshal(), clientSigner.PublicKey().Marshal()) {
				return &ssh.Permissions{}, nil
			}
			return nil, io.EOF
		},
	}
	cfg.AddHostKey(hostSigner)

	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("listen: %v", err)
	}
	t.Cleanup(func() { ln.Close() })
	srv := &testServer{addr: ln.Addr().String()}

	go func() {
		conn, err := ln.Accept()
		if err != nil {
			return
		}
		sconn, chans, reqs, err := ssh.NewServerConn(conn, cfg)
		if err != nil {
			return // e.g. the wrong-host-key client aborts the handshake
		}
		defer sconn.Close()
		go ssh.DiscardRequests(reqs)
		for newCh := range chans {
			if newCh.ChannelType() != "session" {
				newCh.Reject(ssh.UnknownChannelType, "unsupported")
				continue
			}
			ch, chReqs, err := newCh.Accept()
			if err != nil {
				continue
			}
			go func() {
				for req := range chReqs {
					if req.Type == "exec" {
						var p struct{ Command string }
						ssh.Unmarshal(req.Payload, &p)
						srv.mu.Lock()
						srv.cmd = p.Command
						srv.mu.Unlock()
						req.Reply(true, nil)
						data, _ := io.ReadAll(ch) // the pushed payload (client EOFs stdin)
						srv.mu.Lock()
						srv.captured = data
						srv.mu.Unlock()
						if stderrResp != "" {
							ch.Stderr().Write([]byte(stderrResp))
						}
						if stdoutResp != "" {
							ch.Write([]byte(stdoutResp))
						}
						ch.SendRequest("exit-status", false, ssh.Marshal(struct{ Status uint32 }{exitStatus}))
						ch.Close()
						return
					}
					req.Reply(false, nil)
				}
			}()
		}
	}()
	return srv
}

func hostKeyLine(t *testing.T, s ssh.Signer) string {
	t.Helper()
	return strings.TrimSpace(string(ssh.MarshalAuthorizedKey(s.PublicKey())))
}

func newTestClient(t *testing.T, addr, hostKey string, clientPEM []byte) *Client {
	t.Helper()
	c, err := New(Config{
		Addr: addr, User: "felhom-peersync", PrivateKey: clientPEM,
		HostKeyLine: hostKey, Timeout: 5 * time.Second,
	}, log.New(io.Discard, "", 0))
	if err != nil {
		t.Fatalf("wgsync.New: %v", err)
	}
	return c
}

func TestPush_PayloadDeliveredExactly(t *testing.T) {
	clientPEM, clientSigner := testKeys(t)
	_, hostSigner := testKeys(t)
	srv := startTestServer(t, hostSigner, clientSigner, `{"status":"ok","applied":1}`, "", 0)
	c := newTestClient(t, srv.addr, hostKeyLine(t, hostSigner), clientPEM)

	payload := []byte(`{"version":1,"interface":"wg0","peers":[{"pubkey":"PK1","allowed_ip":"10.77.0.2/32"}]}`)
	if err := c.Push(context.Background(), payload); err != nil {
		t.Fatalf("Push: %v", err)
	}
	srv.mu.Lock()
	defer srv.mu.Unlock()
	if !bytes.Equal(srv.captured, payload) {
		t.Errorf("server captured %q, want the exact payload %q", srv.captured, payload)
	}
	if srv.cmd != "felhom-peersync" {
		t.Errorf("exec command = %q, want felhom-peersync (documents intent even under forced command)", srv.cmd)
	}
}

func TestPush_RemoteFailureSurfacesStderr(t *testing.T) {
	clientPEM, clientSigner := testKeys(t)
	_, hostSigner := testKeys(t)
	srv := startTestServer(t, hostSigner, clientSigner, "", "boom", 1)
	c := newTestClient(t, srv.addr, hostKeyLine(t, hostSigner), clientPEM)

	err := c.Push(context.Background(), []byte(`{}`))
	if err == nil {
		t.Fatal("Push succeeded against exit-1 server")
	}
	if !strings.Contains(err.Error(), "boom") {
		t.Errorf("error %q does not carry the remote stderr", err)
	}
}

func TestPush_MalformedResponseIsError(t *testing.T) {
	clientPEM, clientSigner := testKeys(t)
	_, hostSigner := testKeys(t)
	srv := startTestServer(t, hostSigner, clientSigner, "garbage-not-json", "", 0)
	c := newTestClient(t, srv.addr, hostKeyLine(t, hostSigner), clientPEM)

	if err := c.Push(context.Background(), []byte(`{}`)); err == nil {
		t.Fatal("Push accepted a malformed script response — it must not guess")
	}
}

func TestPush_WrongHostKeyRefused(t *testing.T) {
	clientPEM, clientSigner := testKeys(t)
	_, hostSigner := testKeys(t)
	_, otherSigner := testKeys(t) // NOT the server's key
	srv := startTestServer(t, hostSigner, clientSigner, `{"status":"ok","applied":0}`, "", 0)
	c := newTestClient(t, srv.addr, hostKeyLine(t, otherSigner), clientPEM)

	err := c.Push(context.Background(), []byte(`{}`))
	if err == nil {
		t.Fatal("Push succeeded against a server with the WRONG host key — the pin is dead")
	}
	if !strings.Contains(err.Error(), "host key") && !strings.Contains(err.Error(), "handshake") {
		t.Errorf("error %q is not a host-key refusal", err)
	}
	srv.mu.Lock()
	defer srv.mu.Unlock()
	if len(srv.captured) != 0 {
		t.Errorf("payload leaked to a mis-keyed server: %q", srv.captured)
	}
}

// TestPush_MultiHostkeyServerStillMatchesPin reproduces the live S1 failure: a server holding
// MULTIPLE host keys (like a stock sshd: ECDSA + ed25519) must still present the PINNED type.
// Without constraining ClientConfig.HostKeyAlgorithms to the pinned key's algorithm, default
// negotiation can select the other key and FixedHostKey refuses a legitimate server.
func TestPush_MultiHostkeyServerStillMatchesPin(t *testing.T) {
	clientPEM, clientSigner := testKeys(t)
	_, edSigner := testKeys(t)

	ecKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		t.Fatalf("ecdsa: %v", err)
	}
	ecSigner, err := ssh.NewSignerFromKey(ecKey)
	if err != nil {
		t.Fatalf("ecdsa signer: %v", err)
	}

	// Server holds BOTH keys, ECDSA added first (the tempting default pick).
	cfg := &ssh.ServerConfig{
		PublicKeyCallback: func(conn ssh.ConnMetadata, key ssh.PublicKey) (*ssh.Permissions, error) {
			if bytes.Equal(key.Marshal(), clientSigner.PublicKey().Marshal()) {
				return &ssh.Permissions{}, nil
			}
			return nil, io.EOF
		},
	}
	cfg.AddHostKey(ecSigner)
	cfg.AddHostKey(edSigner)

	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("listen: %v", err)
	}
	t.Cleanup(func() { ln.Close() })
	go func() {
		conn, err := ln.Accept()
		if err != nil {
			return
		}
		sconn, chans, reqs, err := ssh.NewServerConn(conn, cfg)
		if err != nil {
			return
		}
		defer sconn.Close()
		go ssh.DiscardRequests(reqs)
		for newCh := range chans {
			ch, chReqs, err := newCh.Accept()
			if err != nil {
				continue
			}
			go func() {
				for req := range chReqs {
					if req.Type == "exec" {
						req.Reply(true, nil)
						io.ReadAll(ch)
						ch.Write([]byte(`{"status":"ok","applied":0}`))
						ch.SendRequest("exit-status", false, ssh.Marshal(struct{ Status uint32 }{0}))
						ch.Close()
						return
					}
					req.Reply(false, nil)
				}
			}()
		}
	}()

	// Pin the ed25519 key — Push must negotiate exactly that type and succeed.
	c := newTestClient(t, ln.Addr().String(), hostKeyLine(t, edSigner), clientPEM)
	if err := c.Push(context.Background(), []byte(`{"peers":[]}`)); err != nil {
		t.Fatalf("Push against multi-hostkey server with correct pin: %v", err)
	}
}

func TestNew_BadInputsFailEarly(t *testing.T) {
	clientPEM, _ := testKeys(t)
	_, hostSigner := testKeys(t)
	hk := hostKeyLine(t, hostSigner)
	logger := log.New(io.Discard, "", 0)

	if _, err := New(Config{Addr: "x:22", User: "u", PrivateKey: []byte("not-a-key"), HostKeyLine: hk}, logger); err == nil {
		t.Error("bad private key accepted")
	}
	if _, err := New(Config{Addr: "x:22", User: "u", PrivateKey: clientPEM, HostKeyLine: "not a key line"}, logger); err == nil {
		t.Error("bad host key line accepted")
	}
	if _, err := New(Config{User: "u", PrivateKey: clientPEM, HostKeyLine: hk}, logger); err == nil {
		t.Error("missing addr accepted")
	}
}