// Package tunnel provides a reverse SSH tunnel over TLS:443 via a Zoraxy edge proxy.
// It opens an outbound TLS connection to the configured jumphost, authenticates
// via SSH, and establishes a reverse port forward so you can reach the target VM
// through sechpoint.app.
package tunnel

import (
	"crypto/tls"
	"fmt"
	"io"
	"net"
	"os"
	"os/signal"
	"time"

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

// Config holds the tunnel connection parameters.
type Config struct {
	Jumphost     string // e.g., "ssh.sechpoint.app:443"
	RemotePort   int    // Port on the jumphost that forwards to target's SSH
	LocalSSHPort int    // SSH port on the target VM (usually 22)
	User         string // SSH user on the jumphost
	Password     string // Password auth (takes lowest priority)
	KeyPath      string // Path to private key for authentication
	KeyBytes     []byte // Raw private key bytes (takes precedence over KeyPath)
}

// DefaultConfig returns a Config with sensible defaults.
func DefaultConfig() Config {
	return Config{
		Jumphost:     "ssh.sechpoint.app:443",
		RemotePort:   9042,
		LocalSSHPort: 22,
		User:         "wallarm-tunnel",
	}
}

// Start opens a reverse SSH tunnel over TLS and keeps it alive.
// It blocks until SIGINT or connection failure.
func Start(cfg Config) error {
	// Build auth methods
	var authMethods []ssh.AuthMethod

	if len(cfg.KeyBytes) > 0 {
		signer, err := ssh.ParsePrivateKey(cfg.KeyBytes)
		if err == nil {
			authMethods = append(authMethods, ssh.PublicKeys(signer))
		}
	}
	if cfg.KeyPath != "" {
		keyBytes, err := os.ReadFile(cfg.KeyPath)
		if err == nil {
			signer, err := ssh.ParsePrivateKey(keyBytes)
			if err == nil {
				authMethods = append(authMethods, ssh.PublicKeys(signer))
			}
		}
	}
	if cfg.Password != "" {
		authMethods = append(authMethods, ssh.Password(cfg.Password))
	}

	if len(authMethods) == 0 {
		return fmt.Errorf("no authentication method configured (provide key or password)")
	}

	sshConfig := &ssh.ClientConfig{
		User:            cfg.User,
		Auth:            authMethods,
		HostKeyCallback: ssh.InsecureIgnoreHostKey(), // trusted infrastructure
		Timeout:         10 * time.Second,
	}

	// TLS dial to the Zoraxy edge (port 443)
	tlsConn, err := tls.Dial("tcp", cfg.Jumphost, &tls.Config{
		InsecureSkipVerify: false,
	})
	if err != nil {
		return fmt.Errorf("TLS dial %s: %w", cfg.Jumphost, err)
	}

	// SSH over TLS
	sshConn, chans, reqs, err := ssh.NewClientConn(tlsConn, cfg.Jumphost, sshConfig)
	if err != nil {
		tlsConn.Close()
		return fmt.Errorf("SSH handshake: %w", err)
	}
	client := ssh.NewClient(sshConn, chans, reqs)
	defer client.Close()

	// Request reverse port forward: jumphost:RemotePort -> localhost:LocalSSHPort
	remoteAddr := fmt.Sprintf("0.0.0.0:%d", cfg.RemotePort)
	localAddr := fmt.Sprintf("localhost:%d", cfg.LocalSSHPort)

	listener, err := client.Listen("tcp", remoteAddr)
	if err != nil {
		return fmt.Errorf("remote listen %s: %w", remoteAddr, err)
	}
	defer listener.Close()

	fmt.Printf("Tunnel established: %s -> %s\n", remoteAddr, localAddr)
	fmt.Printf("Connect: ssh -p %d root@%s\n", cfg.RemotePort, cfg.Jumphost)

	// Handle incoming connections on the remote listener
	go func() {
		for {
			remoteConn, err := listener.Accept()
			if err != nil {
				return
			}
			go forwardConnection(remoteConn, localAddr)
		}
	}()

	// Keep alive until signal
	sigCh := make(chan os.Signal, 1)
	signal.Notify(sigCh, os.Interrupt)

	// Heartbeat every 30s
	ticker := time.NewTicker(30 * time.Second)
	defer ticker.Stop()

	for {
		select {
		case <-sigCh:
			fmt.Println("\nTunnel closed.")
			return nil
		case <-ticker.C:
			_, _, err := client.SendRequest("keepalive@wallarm", true, nil)
			if err != nil {
				return fmt.Errorf("keepalive failed: %w", err)
			}
		}
	}
}

func forwardConnection(remoteConn net.Conn, localAddr string) {
	defer remoteConn.Close()
	localConn, err := net.DialTimeout("tcp", localAddr, 10*time.Second)
	if err != nil {
		return
	}
	defer localConn.Close()

	go func() {
		io.Copy(localConn, remoteConn)
		remoteConn.Close()
	}()
	io.Copy(remoteConn, localConn)
}