package connector

import (
	"context"
	"fmt"
	"io"
	"net"
	"time"

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

type RunOptions struct {
	Command string
	Stdin   io.Reader
	Stdout  io.Writer
	Stderr  io.Writer
}

type DialFn func(ctx context.Context, info provider.ConnectInfo, options DialOptions) (Client, error)

type DialOptions struct {
	DialFn          func(ctx context.Context, network string, address string) (net.Conn, error)
	UseExternalAddr bool
}

type ConnectorOptions struct {
	RunOptions
	DialOptions
}

//go:generate mockery --name=Client --with-expecter
type Client interface {
	Run(ctx context.Context, opts RunOptions) error
	Dial(network string, address string) (net.Conn, error)
	DialRun(ctx context.Context, command string) (net.Conn, error)
	Close() error
}

type ExitError struct {
	err      error
	exitCode int
}

func (e *ExitError) Error() string {
	if e.err == nil {
		return fmt.Sprintf("exit code %d", e.exitCode)
	}
	return e.err.Error()
}

func (e *ExitError) Unwrap() error {
	return e.err
}

func (e *ExitError) Is(err error) bool {
	_, ok := err.(*ExitError)
	return ok
}

func (e *ExitError) ExitCode() int {
	return e.exitCode
}

var Dial DialFn = func(ctx context.Context, info provider.ConnectInfo, options DialOptions) (Client, error) {
	switch info.Protocol {
	default:
		client, err := DialSSH(ctx, info, options)
		if err != nil {
			return nil, fmt.Errorf("dial ssh: %w", err)
		}
		return client, nil

	case provider.ProtocolWinRM:
		client, err := DialWinRM(ctx, info, options)
		if err != nil {
			return nil, fmt.Errorf("dial winrm: %w", err)
		}
		return client, nil
	}
}

func Run(ctx context.Context, info provider.ConnectInfo, options ConnectorOptions) error {
	client, err := Dial(ctx, info, options.DialOptions)
	if err != nil {
		return err
	}
	defer client.Close()

	return client.Run(ctx, options.RunOptions)
}

// dialer is a retry-dialer. Typically, a connection timeout set by the OS will
// always take precedence over any caller provided timeout. This dialer
// continues to dial after an OS interruption with whatever time it has left.
type dialer struct {
	Timeout   time.Duration
	KeepAlive time.Duration
	DialFn    func(ctx context.Context, network string, address string) (net.Conn, error)
}

func (d *dialer) DialContext(ctx context.Context, network, address string) (conn net.Conn, err error) {
	ctx, cancel := context.WithTimeout(ctx, d.Timeout)
	defer cancel()

	var dialFn func(ctx context.Context, network string, address string) (net.Conn, error)
	if d.DialFn == nil {
		base := net.Dialer{KeepAlive: d.KeepAlive}
		dialFn = base.DialContext
	} else {
		dialFn = d.DialFn
	}

	for ctx.Err() == nil {
		conn, err = dialFn(ctx, network, address)
		if err == nil {
			break
		}
	}
	if err != nil {
		return conn, err
	}

	return conn, ctx.Err()
}

func hostport(host string, defaultPort string) string {
	h, p, err := net.SplitHostPort(host)
	if err != nil {
		return net.JoinHostPort(host, defaultPort)
	}
	if p == "" {
		p = defaultPort
	}

	return net.JoinHostPort(h, p)
}
