// Package agentproxy relays a single TCP connection between a local service and
// the control plane, so a control plane that cannot route to this host's network
// can still open a console session.
//
// The dial host is hardcoded to loopback. The control plane supplies only a
// port, and nothing in this package can be made to dial anywhere else.
package agentproxy

import (
	"errors"
	"fmt"
	"io"
	"net"
	"strconv"
	"time"

	"gitea.hostxtra.co.uk/vantage/vantage-shared/grpc/pb"
)

const (
	loopbackHost = "127.0.0.1"
	chunkSize    = 32 * 1024
	dialTimeout  = 10 * time.Second
)

// Stream is the agent's half of a ProxyStream.
type Stream interface {
	Send(*pb.ProxyClientMsg) error
	Recv() (*pb.ProxyServerMsg, error)
	CloseSend() error
}

// Open dials the local port, announces itself on the stream, and relays until
// either side ends. A refused dial is reported as an explicit close so the
// operator sees a reason rather than a hang.
func Open(stream Stream, serverID, agentToken, proxyID string, port uint32) error {
	conn, dialErr := net.DialTimeout("tcp",
		net.JoinHostPort(loopbackHost, strconv.Itoa(int(port))), dialTimeout)

	if err := stream.Send(&pb.ProxyClientMsg{Open: &pb.ProxyOpen{
		ServerId:   serverID,
		AgentToken: agentToken,
		ProxyId:    proxyID,
	}}); err != nil {
		if conn != nil {
			_ = conn.Close()
		}
		return fmt.Errorf("send open: %w", err)
	}

	if dialErr != nil {
		_ = stream.Send(&pb.ProxyClientMsg{Close: &pb.ProxyClose{
			Reason: "dial_refused: " + dialErr.Error(),
		}})
		_ = stream.CloseSend()
		return fmt.Errorf("dial 127.0.0.1:%d: %w", port, dialErr)
	}
	defer conn.Close()

	return relay(conn, stream)
}

func relay(conn net.Conn, stream Stream) error {
	errCh := make(chan error, 2)

	// local service -> control plane
	go func() {
		buf := make([]byte, chunkSize)
		for {
			n, err := conn.Read(buf)
			if n > 0 {
				chunk := make([]byte, n)
				copy(chunk, buf[:n])
				if sendErr := stream.Send(&pb.ProxyClientMsg{Data: chunk}); sendErr != nil {
					errCh <- sendErr
					return
				}
			}
			if err != nil {
				errCh <- err
				return
			}
		}
	}()

	// control plane -> local service
	go func() {
		for {
			msg, err := stream.Recv()
			if err != nil {
				errCh <- err
				return
			}
			if msg.Close != nil {
				errCh <- fmt.Errorf("server closed relay: %s", msg.Close.Reason)
				return
			}
			if len(msg.Data) > 0 {
				if _, err := conn.Write(msg.Data); err != nil {
					errCh <- err
					return
				}
			}
		}
	}()

	err := <-errCh
	_ = conn.Close()
	_ = stream.CloseSend()
	if errors.Is(err, io.EOF) {
		return nil
	}
	return err
}