package main

import (
	"crypto/ed25519"
	"crypto/rand"
	"errors"
	"fmt"
	"io"
	"net"
	"os/exec"
	"testing"

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

// execHandler runs the remote command for one session and returns its exit
// code.
type execHandler func(cmd string, stdin io.Reader, stdout, stderr io.Writer) int

// testServer is a minimal in-process SSH server for tests: one host key, one
// accepted client public key, and a pluggable exec handler. No sshd needed.
type testServer struct {
	ln      net.Listener
	hostKey ssh.Signer
	authPub ssh.PublicKey
	handle  execHandler
}

func newTestServer(t *testing.T, clientPub ssh.PublicKey, handle execHandler) *testServer {
	t.Helper()
	_, hpriv, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		t.Fatal(err)
	}
	signer, err := ssh.NewSignerFromKey(hpriv)
	if err != nil {
		t.Fatal(err)
	}
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	s := &testServer{ln: ln, hostKey: signer, authPub: clientPub, handle: handle}
	go s.serve()
	t.Cleanup(func() { ln.Close() })
	return s
}

func (s *testServer) port() int { return s.ln.Addr().(*net.TCPAddr).Port }

func (s *testServer) serve() {
	cfg := &ssh.ServerConfig{
		PublicKeyCallback: func(_ ssh.ConnMetadata, key ssh.PublicKey) (*ssh.Permissions, error) {
			if s.authPub != nil && string(key.Marshal()) == string(s.authPub.Marshal()) {
				return &ssh.Permissions{}, nil
			}
			return nil, fmt.Errorf("unknown public key")
		},
	}
	cfg.AddHostKey(s.hostKey)
	for {
		nConn, err := s.ln.Accept()
		if err != nil {
			return
		}
		go s.handleConn(nConn, cfg)
	}
}

func (s *testServer) handleConn(nConn net.Conn, cfg *ssh.ServerConfig) {
	conn, chans, reqs, err := ssh.NewServerConn(nConn, cfg)
	if err != nil {
		return
	}
	defer conn.Close()
	go ssh.DiscardRequests(reqs)
	for newCh := range chans {
		if newCh.ChannelType() != "session" {
			newCh.Reject(ssh.UnknownChannelType, "only session")
			continue
		}
		ch, requests, err := newCh.Accept()
		if err != nil {
			return
		}
		go s.handleSession(ch, requests)
	}
}

func (s *testServer) handleSession(ch ssh.Channel, requests <-chan *ssh.Request) {
	for req := range requests {
		if req.Type != "exec" {
			if req.WantReply {
				req.Reply(false, nil)
			}
			continue
		}
		var payload struct{ Command string }
		ssh.Unmarshal(req.Payload, &payload)
		if req.WantReply {
			req.Reply(true, nil)
		}
		code := s.handle(payload.Command, ch, ch, ch.Stderr())
		ch.SendRequest("exit-status", false, ssh.Marshal(struct{ Status uint32 }{uint32(code)}))
		ch.Close()
		return
	}
}

// shellExec runs cmd via /bin/sh -c, exactly as sshd would, wiring the channel
// to the process. Used by the real-rsync end-to-end test.
//
// stdin is fed through a detached goroutine rather than c.Stdin so that Wait
// does not block on the stdin copy: rsync's remote receiver exits while the
// client still holds the channel open, and a c.Run with a non-file Stdin would
// deadlock waiting for that copy to finish.
func shellExec(cmd string, stdin io.Reader, stdout, stderr io.Writer) int {
	c := exec.Command("/bin/sh", "-c", cmd)
	c.Stdout = stdout
	c.Stderr = stderr
	inPipe, err := c.StdinPipe()
	if err != nil {
		return 1
	}
	if err := c.Start(); err != nil {
		return 1
	}
	go func() {
		io.Copy(inPipe, stdin)
		inPipe.Close()
	}()
	err = c.Wait()
	var ee *exec.ExitError
	if errors.As(err, &ee) {
		return ee.ExitCode()
	}
	if err != nil {
		return 1
	}
	return 0
}