package outbound

import (
	"context"
	"errors"
	"fmt"
	"golang.org/x/crypto/ssh"
	"linkport/internal/model"
	"net"
	"strconv"
	"time"
)

// Each forwarded connection owns its SSH transport, including deadlines and cleanup.
type sshForwardConn struct {
	net.Conn
	raw    net.Conn
	client *ssh.Client
}

func (c *sshForwardConn) Close() error { err := c.Conn.Close(); c.client.Close(); return err }
func (c *sshForwardConn) CloseWrite() error {
	if half, ok := c.Conn.(interface{ CloseWrite() error }); ok {
		return half.CloseWrite()
	}
	return errors.New("SSH stream does not support half-close")
}
func (c *sshForwardConn) SetDeadline(t time.Time) error      { return c.raw.SetDeadline(t) }
func (c *sshForwardConn) SetReadDeadline(t time.Time) error  { return c.raw.SetReadDeadline(t) }
func (c *sshForwardConn) SetWriteDeadline(t time.Time) error { return c.raw.SetWriteDeadline(t) }

func dialSSH(ctx context.Context, network, address string, o model.Outbound, keepAlive ...time.Duration) (net.Conn, error) {
	if network != "tcp" {
		return nil, errors.New("SSH forwarding requires TCP")
	}
	client, raw, err := connectSSH(ctx, o, keepAlive...)
	if err != nil {
		return nil, err
	}
	stop := context.AfterFunc(ctx, func() { raw.Close() })
	defer stop()
	done := Begin(ctx, "ssh_forward_channel")
	forwarded, err := client.DialContext(ctx, "tcp", address)
	done(err)
	if err != nil {
		client.Close()
		return nil, fmt.Errorf("SSH forwarding target connection: %w", err)
	}
	if !stop() || ctx.Err() != nil {
		forwarded.Close()
		client.Close()
		return nil, context.Canceled
	}
	raw.SetDeadline(time.Time{})
	return &sshForwardConn{Conn: forwarded, raw: raw, client: client}, nil
}

// TestSSH authenticates the jump host without opening a target channel.
func TestSSH(ctx context.Context, o model.Outbound) (model.TestResult, error) {
	ctx, cancel := context.WithTimeout(ctx, 15*time.Second)
	defer cancel()
	var fingerprint string
	client, _, err := connectSSHObserved(ctx, o, &fingerprint)
	if err != nil {
		return model.TestResult{}, err
	}
	defer client.Close()
	return model.TestResult{Version: string(client.ServerVersion()), Fingerprint: fingerprint}, nil
}

// ProbeSSH aborts at the host key callback, before any user authentication.
func ProbeSSH(ctx context.Context, host string, port int) (string, error) {
	ctx, cancel := context.WithTimeout(ctx, 15*time.Second)
	defer cancel()
	endpoint := net.JoinHostPort(host, strconv.Itoa(port))
	raw, err := (&net.Dialer{}).DialContext(ctx, "tcp", endpoint)
	if err != nil {
		return "", err
	}
	defer raw.Close()
	stop := context.AfterFunc(ctx, func() { raw.Close() })
	defer stop()
	deadline, _ := ctx.Deadline()
	raw.SetDeadline(deadline)
	observed := ""
	done := errors.New("host key probe complete")
	config := &ssh.ClientConfig{HostKeyCallback: func(_ string, _ net.Addr, key ssh.PublicKey) error {
		observed = ssh.FingerprintSHA256(key)
		return done
	}}
	_, _, _, err = ssh.NewClientConn(raw, endpoint, config)
	if ctx.Err() != nil {
		return "", ctx.Err()
	}
	if !errors.Is(err, done) || observed == "" {
		return "", fmt.Errorf("SSH host key probe: %w", err)
	}
	return observed, nil
}

func connectSSH(ctx context.Context, o model.Outbound, keepAlive ...time.Duration) (*ssh.Client, net.Conn, error) {
	return connectSSHObserved(ctx, o, nil, keepAlive...)
}

func connectSSHObserved(ctx context.Context, o model.Outbound, observed *string, keepAlive ...time.Duration) (*ssh.Client, net.Conn, error) {
	if o.Auth != "password" || (!o.SkipHostKeyVerify && o.HostFingerprint == "") {
		return nil, nil, errors.New("SSH forwarding authentication or host fingerprint missing")
	}
	endpoint := net.JoinHostPort(o.Host, strconv.Itoa(o.Port))
	dialer := net.Dialer{}
	if len(keepAlive) > 0 {
		dialer.KeepAlive = keepAlive[0]
	}
	done := Begin(ctx, "ssh_forward_connection")
	raw, err := dialer.DialContext(ctx, "tcp", endpoint)
	done(err)
	if err != nil {
		return nil, nil, fmt.Errorf("SSH forwarding connection: %w", err)
	}
	good := false
	defer func() {
		if !good {
			raw.Close()
		}
	}()
	stop := context.AfterFunc(ctx, func() { raw.Close() })
	defer stop()
	deadline, _ := ctx.Deadline()
	raw.SetDeadline(deadline)
	// PasswordCallback observes the actual credential attempt. In particular, a
	// rejected host key must leave authentication not_run.
	authStarted := false
	authDone := func(error) {}
	config := &ssh.ClientConfig{User: o.Username, Auth: []ssh.AuthMethod{ssh.PasswordCallback(func() (string, error) {
		if !authStarted {
			authStarted = true
			authDone = Begin(ctx, "ssh_forward_authentication")
		}
		return o.Password, nil
	})}, HostKeyCallback: func(_ string, _ net.Addr, key ssh.PublicKey) error {
		fingerprint := ssh.FingerprintSHA256(key)
		if observed != nil {
			*observed = fingerprint
		}
		if !o.SkipHostKeyVerify {
			identityDone := Begin(ctx, "ssh_forward_identity")
			if fingerprint != o.HostFingerprint {
				err := model.AtConnectionStage("ssh_forward_identity", errors.New("SSH forwarding host fingerprint mismatch"))
				identityDone(err)
				return err
			}
			identityDone(nil)
		}
		return nil
	}}
	done = Begin(ctx, "ssh_forward_handshake")
	conn, channels, requests, err := ssh.NewClientConn(raw, endpoint, config)
	done(err)
	if err == nil && !authStarted {
		authDone = Begin(ctx, "ssh_forward_authentication")
	}
	authDone(err)
	if err != nil {
		return nil, nil, fmt.Errorf("SSH forwarding handshake: %w", err)
	}
	client := ssh.NewClient(conn, channels, requests)
	if !stop() || ctx.Err() != nil {
		client.Close()
		return nil, nil, context.Canceled
	}
	good = true
	return client, raw, nil
}