package main

import (
	"fmt"
	"io"
	"log/slog"
	"net"
	"sync"

	"golang.org/x/crypto/ssh"
)

type PortForwarder struct {
	sshCon     *ssh.Client
	localPort  string
	remotePort string
	remoteHost string
	logMutex   sync.Mutex
}

func NewPortForwarder(sshCon *ssh.Client, localPort, remotePort, remoteHost string) *PortForwarder {
	return &PortForwarder{
		sshCon:     sshCon,
		localPort:  localPort,
		remotePort: remotePort,
		remoteHost: remoteHost,
	}
}

func (pf *PortForwarder) forward() error {
	localAddr := "127.0.0.1:" + pf.localPort
	remoteAddr := net.JoinHostPort(pf.remoteHost, pf.remotePort)

	pf.logf("INFO", "Starting port forwarder: local %s -> remote %s (via SSH)", localAddr, remoteAddr)

	listener, err := net.Listen("tcp", localAddr)
	if err != nil {
		pf.logf("ERROR", "Failed to open local listener on %s: %v", localAddr, err)
		return fmt.Errorf("failed to listen on %s: %w", localAddr, err)
	}
	defer listener.Close()
	pf.logf("INFO", fmt.Sprintf("Listener active on %s", localAddr))

	for {
		localConn, err := listener.Accept()
		if err != nil {
			if opErr, ok := err.(*net.OpError); ok && opErr.Err.Error() == "use of closed network connection" {
				pf.logf("INFO", "Listener on %s closed, stopping forwarder.", localAddr)
				return nil
			}
			pf.logf("ERROR", "Failed to accept incoming connection on %s: %v", localAddr, err)
			continue
		}

		pf.logf("INFO", fmt.Sprintf("Accepted connection from %s on %s", localConn.RemoteAddr(), localAddr))
		go pf.handleConnection(localConn, remoteAddr)
	}
}

func (pf *PortForwarder) handleConnection(localConn net.Conn, remoteAddr string) {
	defer localConn.Close()

	pf.logf("INFO", "Dialing remote host %s via SSH tunnel for %s", remoteAddr, localConn.RemoteAddr())
	remoteConn, err := pf.sshCon.Dial("tcp", remoteAddr)
	if err != nil {
		pf.logf("ERROR", "Failed to dial remote host %s via SSH: %v", remoteAddr, err)
		return
	}
	defer remoteConn.Close()
	pf.logf("INFO", "Connection to %s established. Starting data copy.", remoteAddr)

	var wg sync.WaitGroup
	wg.Add(2)

	go func() {
		defer wg.Done()
		defer localConn.Close()
		bytesCopied, err := io.Copy(localConn, remoteConn)
		if err != nil {
		}
		pf.logf("INFO", "Finished copying remote->local (%d bytes) for %s", bytesCopied, localConn.RemoteAddr())
	}()

	go func() {
		defer wg.Done()
		defer remoteConn.Close()
		bytesCopied, err := io.Copy(remoteConn, localConn)
		if err != nil {
		}
		pf.logf("INFO", "Finished copying local->remote (%d bytes) for %s", bytesCopied, localConn.RemoteAddr())
	}()

	wg.Wait()
	pf.logf("INFO", "Closing forwarded connection for %s", localConn.RemoteAddr())
}

func (pf *PortForwarder) logf(level, format string, v ...any) {
	pf.logMutex.Lock()
	defer pf.logMutex.Unlock()
	switch level {
	case "INFO":
		slog.Info(format, v...)
	case "WARN":
		slog.Warn(format, v...)
	case "ERROR":
		slog.Error(format, v...)
	}
}