package main

import (
	"fmt"
	"io"
	"log"
	"net"
	"sync" // Mutex hinzufügen für sichereres Logging

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

type PortForwarder struct {
	sshCon     *ssh.Client
	localPort  string
	remotePort string
	remoteHost string
	logMutex   sync.Mutex // Mutex zum Schutz von Log-Ausgaben aus Goroutinen
}

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) // Sicherer Host:Port kombinieren

	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: Listener active on %s", localAddr)

	for {
		localConn, err := listener.Accept()
		if err != nil {
			// Fehler tritt auf, wenn der Listener geschlossen wird oder ein Netzwerkproblem vorliegt.
			// Prüfe, ob der Fehler durch Schließen des Listeners verursacht wurde (erwartet).
			// Fehler wie 'use of closed network connection' sind hier normal beim Beenden.
			// if errors.Is(err, net.ErrClosed) { // Bessere Prüfung in neueren Go-Versionen
			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 // Kein Fehler, normales Beenden
			}
			pf.logf("ERROR: Failed to accept incoming connection on %s: %v", localAddr, err)
			// Optional: Kurze Pause vor erneutem Versuch oder Abbruch?
			// Bei dauerhaften Fehlern wird die Schleife hier schnell laufen.
			// Für dieses Tool ist ein Weiterlaufen bei Accept-Fehlern wahrscheinlich ok.
			continue
		}

		pf.logf("INFO: 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)
		// Schließe lokale Verbindung, wenn Remote nicht erreicht werden kann
		// (defer macht das schon, aber hier explizit zur Klarheit)
		// localConn.Close() // Ist durch defer abgedeckt
		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() // Schließe die eine Seite, wenn die andere endet
		bytesCopied, err := io.Copy(localConn, remoteConn)
		if err != nil {
			// Fehler beim Kopieren sind normal, wenn eine Seite die Verbindung schließt
			// log.Printf("DEBUG: Error copying remote->local: %v", err)
		}
		pf.logf("INFO: Finished copying remote->local (%d bytes) for %s", bytesCopied, localConn.RemoteAddr())
	}()

	go func() {
		defer wg.Done()
		defer remoteConn.Close() // Schließe die andere Seite, wenn diese endet
		bytesCopied, err := io.Copy(remoteConn, localConn)
		if err != nil {
			// log.Printf("DEBUG: Error copying local->remote: %v", err)
		}
		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(format string, v ...interface{}) {
	pf.logMutex.Lock()
	defer pf.logMutex.Unlock()
	log.Printf(format, v...) // Verwende den Standard-Logger
}