package main

import (
	"bytes"
	"context"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"net"
	"strconv"
	"sync"
	"time"
)

var proxySignature = []byte("\r\n\r\n\x00\r\nQUIT\n")

type mailProxy struct {
	listener net.Listener
	upstream string
	source   net.IP
	failures chan<- error
	slots    chan struct{}
	mu       sync.Mutex
	closed   bool
	conns    map[net.Conn]struct{}
	workers  sync.WaitGroup
	done     chan struct{}
}

func newMailProxy(listener net.Listener, upstream string, source net.IP, failures chan<- error, slots chan struct{}) *mailProxy {
	return &mailProxy{listener: listener, upstream: upstream, source: source, failures: failures, slots: slots, conns: make(map[net.Conn]struct{}), done: make(chan struct{})}
}

func (r *relay) startProxies() error {
	if r.config.mailTarget == nil {
		return nil
	}
	slots := make(chan struct{}, mailConnectionLimit)
	for _, port := range r.config.mailTarget.ports {
		listener, err := net.ListenTCP("tcp4", &net.TCPAddr{IP: r.config.mailTarget.address, Port: int(port)})
		if err != nil {
			r.stopProxies()
			return fmt.Errorf("listen on %s:%d: %w", r.config.mailTarget.address, port, err)
		}
		proxy := newMailProxy(listener, net.JoinHostPort(r.config.mailUpstream, strconv.Itoa(int(port))), net.ParseIP(r.config.podIP).To4(), r.proxyFailures, slots)
		r.proxies = append(r.proxies, proxy)
		go proxy.serve()
	}
	return nil
}

func (r *relay) stopProxies() {
	for _, proxy := range r.proxies {
		proxy.close()
	}
	r.proxies = nil
}

func (p *mailProxy) serve() {
	defer close(p.done)
	for {
		client, err := p.listener.Accept()
		if err != nil {
			if !p.isClosed() {
				p.report(fmt.Errorf("accept connection: %w", err))
			}
			return
		}
		select {
		case p.slots <- struct{}{}:
		default:
			_ = client.Close()
			continue
		}
		if !p.track(client) {
			<-p.slots
			_ = client.Close()
			return
		}
		p.workers.Add(1)
		go p.forward(client)
	}
}

func (p *mailProxy) forward(client net.Conn) {
	defer p.workers.Done()
	defer func() { <-p.slots }()
	defer p.untrack(client)
	defer client.Close()
	dialer := net.Dialer{Timeout: 10 * time.Second, KeepAlive: 30 * time.Second}
	if p.source != nil {
		dialer.LocalAddr = &net.TCPAddr{IP: p.source}
	}
	server, err := dialer.DialContext(context.Background(), "tcp4", p.upstream)
	if err != nil {
		return
	}
	defer server.Close()
	if !p.track(server) {
		return
	}
	defer p.untrack(server)
	if err := writeProxyHeader(server, client.RemoteAddr(), client.LocalAddr()); err != nil {
		return
	}
	clientTCP, clientOK := client.(*net.TCPConn)
	serverTCP, serverOK := server.(*net.TCPConn)
	if !clientOK || !serverOK {
		return
	}
	results := make(chan error, 2)
	go copyStream(results, serverTCP, clientTCP)
	go copyStream(results, clientTCP, serverTCP)
	<-results
	<-results
}

func (p *mailProxy) close() {
	p.mu.Lock()
	if p.closed {
		p.mu.Unlock()
		return
	}
	p.closed = true
	connections := make([]net.Conn, 0, len(p.conns))
	for connection := range p.conns {
		connections = append(connections, connection)
	}
	p.mu.Unlock()
	_ = p.listener.Close()
	<-p.done
	for _, connection := range connections {
		_ = connection.Close()
	}
	p.workers.Wait()
}

func (p *mailProxy) track(connection net.Conn) bool {
	p.mu.Lock()
	defer p.mu.Unlock()
	if p.closed {
		return false
	}
	p.conns[connection] = struct{}{}
	return true
}

func (p *mailProxy) untrack(connection net.Conn) {
	p.mu.Lock()
	delete(p.conns, connection)
	p.mu.Unlock()
}

func (p *mailProxy) isClosed() bool {
	p.mu.Lock()
	defer p.mu.Unlock()
	return p.closed
}

func (p *mailProxy) report(err error) {
	select {
	case p.failures <- err:
	default:
	}
}

func copyStream(results chan<- error, destination, source *net.TCPConn) {
	_, err := io.Copy(destination, source)
	if closeErr := destination.CloseWrite(); err == nil {
		err = closeErr
	}
	results <- err
}

func writeProxyHeader(writer io.Writer, sourceAddress, destinationAddress net.Addr) error {
	source, sourceOK := sourceAddress.(*net.TCPAddr)
	destination, destinationOK := destinationAddress.(*net.TCPAddr)
	if !sourceOK || !destinationOK {
		return errors.New("PROXY addresses are not TCP")
	}
	sourceIP := source.IP.To4()
	destinationIP := destination.IP.To4()
	if sourceIP == nil || destinationIP == nil || source.Port < 1 || source.Port > 65535 || destination.Port < 1 || destination.Port > 65535 {
		return errors.New("PROXY addresses are not IPv4")
	}
	header := make([]byte, 28)
	copy(header, proxySignature)
	header[12] = 0x21
	header[13] = 0x11
	binary.BigEndian.PutUint16(header[14:16], 12)
	copy(header[16:20], sourceIP)
	copy(header[20:24], destinationIP)
	binary.BigEndian.PutUint16(header[24:26], uint16(source.Port))
	binary.BigEndian.PutUint16(header[26:28], uint16(destination.Port))
	_, err := io.Copy(writer, bytes.NewReader(header))
	return err
}

func drainFailures(failures <-chan error) {
	for {
		select {
		case <-failures:
		default:
			return
		}
	}
}