package sshtun

import (
	"fmt"
	"net"
	"strconv"
	"strings"

	"github.com/Neur0toxine/sshpoke/pkg/errtools"
	"github.com/Neur0toxine/sshpoke/pkg/proto/ssh"
	"github.com/function61/gokit/io/bidipipe"
	"go.uber.org/zap"
)

type Forward struct {
	// local service to be forwarded
	Local Endpoint `json:"local"`
	// remote forwarding port (on remote SSH server network)
	Remote Endpoint `json:"remote"`
}

func AddrToEndpoint(address string) Endpoint {
	host, port, err := net.SplitHostPort(address)
	if err != nil && errtools.IsPortMissingErr(err) {
		return Endpoint{Host: address, Port: 22}
	}
	portNum, err := strconv.Atoi(port)
	if err != nil {
		portNum = 22
	}
	return Endpoint{Host: host, Port: portNum}
}

type Endpoint struct {
	Host string `json:"host"`
	Port int    `json:"port"`
}

func (endpoint *Endpoint) String() string {
	return fmt.Sprintf("%s:%d", endpoint.Host, endpoint.Port)
}

func (t *Tunnel) ipFromAddr(addr net.Addr) net.IP {
	host, _, _ := net.SplitHostPort(addr.String())
	return net.ParseIP(host)
}

//	blocking flow: calls Listen() on the SSH connection, and if succeeds returns non-nil error
//
// nonblocking flow: if Accept() call fails, stops goroutine and returns error on ch listenerStopped
func (t *Tunnel) reverseForwardOnePort(sshClient *ssh.Client, listenerStopped chan<- error) error {
	var (
		listener net.Listener
		err      error
	)
	if t.tunConfig.FakeRemoteHost {
		listener, err = sshClient.ListenTCP(&net.TCPAddr{
			IP:   t.ipFromAddr(sshClient.Conn.RemoteAddr()),
			Port: t.tunConfig.Forward.Remote.Port,
		}, t.tunConfig.Forward.Remote.Host)
	} else {
		listener, err = sshClient.Listen("tcp", t.tunConfig.Forward.Remote.String())
	}
	if err != nil {
		return err
	}

	go func() {
		defer listener.Close()
		t.log.Debugf("forwarding %s <- %s",
			t.tunConfig.Forward.Local.String(), t.tunConfig.Forward.Remote.String())

		for {
			client, err := listener.Accept()
			if err != nil {
				listenerStopped <- fmt.Errorf("error on Accept(): %w", err)
				return
			}

			go handleReverseForwardConn(client, t.tunConfig.Forward, t.log)
		}
	}()

	return nil
}

func handleReverseForwardConn(client net.Conn, forward Forward, log *zap.SugaredLogger) {
	defer client.Close()

	remote, err := net.Dial("tcp", forward.Local.String())
	if err != nil {
		log.Errorf("cannot dial local service: %s", err.Error())
		return
	}

	log.Debugf("proxying %s <-> %s", forward.Local.String(), client.RemoteAddr())

	// pipe data in both directions:
	// - client => remote
	// - remote => client
	//
	// - in effect, we act as a proxy between the reverse tunnel's client and locally-dialed
	//   remote endpoint.
	// - the "client" and "remote" strings we give Pipe() is just for error&log messages
	// - this blocks until either of the parties' socket closes (or breaks)
	if err := bidipipe.Pipe(
		bidipipe.WithName("client", client),
		bidipipe.WithName("remote", remote),
	); err != nil {
		// we can safely ignore those errors.
		if strings.Contains(err.Error(), "use of closed network connection") {
			return
		}
		log.Error(err)
	}
}