package connector

import (
	"context"
	"errors"
	"fmt"
	"io"
	"net"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"golang.org/x/crypto/ssh"

	"gitlab.com/gitlab-org/fleeting/fleeting/provider"
)

type sshClient struct {
	client *ssh.Client
	done   chan struct{}

	mu     sync.Mutex
	closed bool
}

const (
	// handshakeErrorRetries is the number of retries performed specifically
	// for "handshake"/initial temporary connection related errors.
	//
	// These occur often when a new instance is being created because the SSH
	// socket might be available, but setup via init scripts are still
	// configuring and/or restarting the daemon.
	//
	// 30 attempts and at a 2-second interval is pretty long for this
	// specific type of problem, however, it's also capped by the context
	// deadline and the timeout provided for it. So it being too long is
	// better than it being too short.
	handshakeErrorRetries       = 30
	handshakeErrorRetryInterval = 2 * time.Second
)

func DialSSH(ctx context.Context, info provider.ConnectInfo, options DialOptions) (*sshClient, error) {
	addr := info.InternalAddr
	if options.UseExternalAddr && (info.ExternalAddr != "" || addr == "") {
		addr = info.ExternalAddr
	}
	addr = hostport(addr, "22")

	config := &ssh.ClientConfig{
		User:            info.Username,
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
		Timeout:         info.Timeout,
	}

	if info.Key == nil {
		config.Auth = append(config.Auth, ssh.Password(info.Password))
	} else {
		signer, err := ssh.ParsePrivateKey(info.Key)
		if err != nil {
			return nil, err
		}
		config.Auth = append(config.Auth, ssh.PublicKeys(signer))
	}

	dialer := dialer{
		Timeout:   info.Timeout,
		KeepAlive: info.Keepalive,
		DialFn:    options.DialFn,
	}

	ioTimeout := info.Timeout

	var client *ssh.Client
	var err error
	i := 0
	for ; i < handshakeErrorRetries; i++ {
		startedDial := time.Now()
		client, err = getSSHClient(ctx, dialer, addr, config, ioTimeout)

		// the io timeout is shared between retries, so we remove how long
		// the last attempt took
		ioTimeout -= time.Since(startedDial)

		if err != nil && isTemporarySSHError(err) && ioTimeout > 0 {
			time.Sleep(handshakeErrorRetryInterval)
			continue
		}
		break
	}
	if err != nil {
		return nil, fmt.Errorf("after retrying %d times: %w", i, err)
	}

	sc := &sshClient{client: client, done: make(chan struct{})}
	go sc.keepalive(info.Keepalive)

	return sc, nil
}

func isTemporarySSHError(err error) bool {
	str := err.Error()

	return strings.Contains(str, "handshake failed") || strings.Contains(str, "i/o timeout")
}

func (c *sshClient) keepalive(interval time.Duration) {
	for {
		select {
		case <-c.done:
			return
		default:
			time.Sleep(interval)
			_, _, _ = c.client.SendRequest("keepalive@openssh.com", true, nil)
		}
	}
}

func (c *sshClient) Close() error {
	c.mu.Lock()
	defer c.mu.Unlock()

	if c.closed {
		return fmt.Errorf("connection already closed")
	}

	c.closed = true
	close(c.done)

	return c.client.Close()
}

func (c *sshClient) Run(ctx context.Context, opts RunOptions) error {
	sess, err := c.client.NewSession()
	if err != nil {
		return err
	}
	defer sess.Close()

	var stderrWriter *stderrOmitWriter
	if opts.Stderr == nil {
		stderrWriter = newStderrOmitWriter()
		opts.Stderr = stderrWriter
	}

	sess.Stdin = opts.Stdin
	sess.Stdout = opts.Stdout
	sess.Stderr = opts.Stderr

	errCh := make(chan error, 1)

	go func() {
		errCh <- sess.Run(opts.Command)
	}()

	select {
	case <-ctx.Done():
		return sess.Close()
	case err := <-errCh:
		if err != nil && stderrWriter != nil {
			err = fmt.Errorf("%w (%w)", err, stderrWriter.Error())
		}

		return c.newExitError(err)
	}
}

func (c *sshClient) newExitError(err error) error {
	if err == nil {
		return nil
	}

	var sshExitErr *ssh.ExitError
	if errors.As(err, &sshExitErr) {
		return &ExitError{
			err:      err,
			exitCode: sshExitErr.ExitStatus(),
		}
	}

	return &ExitError{err: err, exitCode: 0}
}

func (c *sshClient) DialRun(ctx context.Context, command string) (net.Conn, error) {
	connReader, w := io.Pipe()
	r, connWriter := io.Pipe()

	go func() {
		ctx, cancel := context.WithCancel(ctx)
		defer cancel()

		defer r.Close()
		defer w.Close()

		stderr := newStderrOmitWriter()
		err := c.Run(ctx, RunOptions{
			Command: command,
			Stdin:   r,
			Stdout:  w,
			Stderr:  stderr,
		})
		if err != nil {
			err = fmt.Errorf("ssh tunnel: %w (%s)", err, stderr.Error())
			r.CloseWithError(err)
			w.CloseWithError(err)
		}
	}()

	return &rwConn{connWriter, connReader}, nil
}

func (c *sshClient) Dial(network string, address string) (net.Conn, error) {
	return c.client.Dial(network, address)
}

func getSSHClient(ctx context.Context, d dialer, address string, config *ssh.ClientConfig, ioTimeout time.Duration) (*ssh.Client, error) {
	ctx, cancel := context.WithTimeout(ctx, ioTimeout)
	defer cancel()

	conn, err := d.DialContext(ctx, "tcp", address)
	if err != nil {
		return nil, err
	}

	// wait until the context is canceled:
	// if success is true, we don't close the connection
	// if success if false, we close the connection as it stops NewClientConn
	// from continuing to block.
	var success atomic.Bool
	go func() {
		<-ctx.Done()
		if !success.Load() {
			conn.Close()
		}
	}()

	defer conn.SetDeadline(time.Time{})

	// setting dial deadline is best effort, if the context is canceled,
	// we close the entire connection anyway.
	_ = conn.SetDeadline(time.Now().Add(ioTimeout))

	c, chans, reqs, err := ssh.NewClientConn(conn, address, config)
	if err != nil {
		return nil, err
	}

	success.Store(true)

	return ssh.NewClient(c, chans, reqs), nil
}
