package main

import (
	"context"
	"flag"
	"fmt"
	"io"
	"log/slog"
	"net"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"
)

func main() {
	listenAddress := flag.String("listen", "127.0.0.1:10554", "local listen address")
	upstreamAddress := flag.String("upstream", "", "upstream host:port")
	flag.Parse()
	logger := slog.New(slog.NewJSONHandler(os.Stdout, nil))
	if *upstreamAddress == "" {
		logger.Error("upstream is required")
		os.Exit(2)
	}
	if _, _, err := net.SplitHostPort(*upstreamAddress); err != nil {
		logger.Error("upstream must be host:port")
		os.Exit(2)
	}
	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer stop()
	if err := serve(ctx, *listenAddress, *upstreamAddress); err != nil {
		logger.Error("RTSP fault proxy stopped", "error", err)
		os.Exit(1)
	}
}

func serve(ctx context.Context, listenAddress, upstreamAddress string) error {
	listener, err := net.Listen("tcp", listenAddress)
	if err != nil {
		return fmt.Errorf("listen: %w", err)
	}
	defer listener.Close()
	go func() {
		<-ctx.Done()
		_ = listener.Close()
	}()
	var connections sync.WaitGroup
	defer connections.Wait()
	for {
		client, acceptErr := listener.Accept()
		if acceptErr != nil {
			if ctx.Err() != nil {
				return nil
			}
			return fmt.Errorf("accept: %w", acceptErr)
		}
		connections.Add(1)
		go func() {
			defer connections.Done()
			proxy(client, upstreamAddress)
		}()
	}
}

func proxy(client net.Conn, upstreamAddress string) {
	defer client.Close()
	upstream, err := net.DialTimeout("tcp", upstreamAddress, 5*time.Second)
	if err != nil {
		return
	}
	defer upstream.Close()
	done := make(chan struct{}, 2)
	copyOneWay := func(destination, source net.Conn) {
		_, _ = io.Copy(destination, source)
		done <- struct{}{}
	}
	go copyOneWay(upstream, client)
	go copyOneWay(client, upstream)
	<-done
}