// Package proxy transports an opaque bidirectional byte stream.
package proxy

import (
	"errors"
	"io"
	"net"
)

type copyResult struct {
	err error
	dst net.Conn
}

type closeWriter interface {
	CloseWrite() error
}

// TrafficObserver receives each positive byte count as it is copied.
// clientToBackend identifies the PostgreSQL client-to-server direction.
type TrafficObserver func(clientToBackend bool, bytes int64)

// Bidirectional copies bytes in both directions. TCP peers retain half-close
// semantics; a TLS client is closed when its upstream direction has ended,
// because crypto/tls exposes no CloseWrite operation.
func Bidirectional(client, backend net.Conn) error {
	return BidirectionalWithTraffic(client, backend, nil)
}

// BidirectionalWithTraffic copies bytes in both directions and reports every
// positive read to observer when one is provided.
func BidirectionalWithTraffic(client, backend net.Conn, observer TrafficObserver) error {
	results := make(chan copyResult, 2)
	copyStream := func(dst, src net.Conn, clientToBackend bool) {
		_, err := io.Copy(dst, trafficReader{Reader: src, clientToBackend: clientToBackend, observer: observer})
		results <- copyResult{err: err, dst: dst}
	}
	go copyStream(backend, client, true)
	go copyStream(client, backend, false)

	first := <-results
	if writer, ok := first.dst.(closeWriter); ok {
		_ = writer.CloseWrite()
	} else {
		// TLS cannot be half-closed. Closing unblocks the opposite copy and
		// prevents an idle peer from retaining a handler goroutine.
		_ = client.Close()
	}
	second := <-results
	if writer, ok := second.dst.(closeWriter); ok {
		_ = writer.CloseWrite()
	}

	return combine(first.err, second.err)
}

type trafficReader struct {
	io.Reader
	clientToBackend bool
	observer        TrafficObserver
}

func (r trafficReader) Read(buffer []byte) (int, error) {
	bytes, err := r.Reader.Read(buffer)
	if bytes > 0 && r.observer != nil {
		r.observer(r.clientToBackend, int64(bytes))
	}
	return bytes, err
}

func combine(errs ...error) error {
	var relevant []error
	for _, err := range errs {
		if err != nil && !errors.Is(err, io.EOF) && !errors.Is(err, net.ErrClosed) {
			relevant = append(relevant, err)
		}
	}
	return errors.Join(relevant...)
}

// Compile-time assertion: a TCP connection has practical half-close support.
var _ closeWriter = (*net.TCPConn)(nil)