package connector

import (
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"net"
	"strings"
	"sync"

	"github.com/masterzen/winrm"

	"gitlab.com/gitlab-org/fleeting/fleeting/connector/internal/proxy"
	"gitlab.com/gitlab-org/fleeting/fleeting/provider"
)

type winRMClient struct {
	client *winrm.Client
	info   provider.ConnectInfo

	mu sync.Mutex
}

func DialWinRM(ctx context.Context, info provider.ConnectInfo, opts DialOptions) (*winRMClient, error) {
	addr := info.InternalAddr
	if opts.UseExternalAddr && (info.ExternalAddr != "" || addr == "") {
		addr = info.ExternalAddr
	}

	host, sport, err := net.SplitHostPort(hostport(addr, "5985"))
	if err != nil {
		return nil, fmt.Errorf("parsing address: %w", err)
	}

	port, err := net.LookupPort("tcp", sport)
	if err != nil {
		return nil, fmt.Errorf("resolving port: %w", err)
	}

	endpoint := winrm.NewEndpoint(host, port, false, false, nil, nil, nil, info.Timeout)

	client, err := winrm.NewClientWithParameters(
		winrm.NewEndpoint(host, port, false, false, nil, nil, nil, info.Timeout),
		info.Username,
		info.Password,
		&winrm.Parameters{
			Timeout:      "PT60S",
			EnvelopeSize: 5 * 1024 * 1024,
			TransportDecorator: func() winrm.Transporter {
				return winrm.NewClientWithDial(func(network, addr string) (net.Conn, error) {
					d := dialer{
						Timeout:   endpoint.Timeout,
						KeepAlive: info.Keepalive,
						DialFn:    opts.DialFn,
					}

					return d.DialContext(context.Background(), network, addr)
				})
			},
		},
	)
	if err != nil {
		return nil, fmt.Errorf("creating client: %w", err)
	}

	return &winRMClient{client: client, info: info}, nil
}

func (c *winRMClient) Close() error {
	return nil
}

func (c *winRMClient) Run(ctx context.Context, opts RunOptions) error {
	var stderrWriter *stderrOmitWriter
	if opts.Stderr == nil {
		stderrWriter = newStderrOmitWriter()
		opts.Stderr = stderrWriter
	}

	exitCode, err := c.runWinRMCmd(ctx, opts.Command, opts.Stdout, opts.Stderr, opts.Stdin)
	if exitCode == 0 && err == nil {
		return nil
	}

	if err != nil && stderrWriter != nil {
		err = fmt.Errorf("%w (%w)", err, stderrWriter.Error())
	}

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

func (c *winRMClient) 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.runWinRMCmd(ctx, command, w, stderr, r)
		if err != nil {
			err = fmt.Errorf("winrm tunnel: %w (%s)", err, stderr.Error())
			r.CloseWithError(err)
			w.CloseWithError(err)
		}
	}()

	return &rwConn{connWriter, connReader}, nil
}

func (c *winRMClient) Dial(network string, address string) (net.Conn, error) {
	version, err := c.deployTunnelBinary()
	if err != nil {
		return nil, fmt.Errorf("deploying fleeting proxy: %w", err)
	}

	return c.DialRun(context.Background(), fmt.Sprintf("fleeting-proxy-%s.exe %s %s", version, network, address))
}

func (c *winRMClient) deployTunnelBinary() (string, error) {
	bin := proxy.ReadBinary(c.info.OS, c.info.Arch)
	if bin == nil {
		return "", fmt.Errorf("found no suitable tunnel binary (os:%s, arch:%s)", c.info.OS, c.info.Arch)
	}

	checksum := sha256.Sum256(bin)
	version := hex.EncodeToString(checksum[:8])

	c.mu.Lock()
	defer c.mu.Unlock()

	const cmd = "powershell.exe -NoProfile -NonInteractive -ExecutionPolicy Bypass -Command -"

	output := new(bytes.Buffer)
	code, err := c.runWinRMCmd(context.Background(), cmd, output, output, strings.NewReader(fmt.Sprintf(proxy.PSCheckTunnelBinary, version)))
	if err == nil && code == 0 {
		return version, nil
	}

	encodedBinary := base64.StdEncoding.EncodeToString(bin)
	code, err = c.runWinRMCmd(context.Background(), cmd, output, output, strings.NewReader(fmt.Sprintf(proxy.PSDeployTunnelBinary, version, encodedBinary)))
	if err != nil || code != 0 {
		return "", fmt.Errorf("deploying tunnel binary: %w (%s)", err, output.String())
	}

	return version, nil
}

func (c *winRMClient) runWinRMCmd(ctx context.Context, command string, stdout, stderr io.Writer, stdin io.Reader) (int, error) {
	code, err := c.client.RunWithContextWithInput(ctx, command, stdout, stderr, stdin)

	var executeCmdErr *winrm.ExecuteCommandError
	if errors.As(err, &executeCmdErr) {
		err = fmt.Errorf("%w: %v", err, executeCmdErr.Body)
	}

	return code, err
}
