//go:build windows

package gcs

import (
	"context"
	"fmt"
	"io"
	"net"
	"strings"
	"sync"

	"go.opentelemetry.io/otel"
	"go.opentelemetry.io/otel/attribute"
	"go.opentelemetry.io/otel/propagation"
	"go.opentelemetry.io/otel/trace"

	"github.com/Microsoft/go-winio"
	"github.com/Microsoft/go-winio/pkg/guid"
	"github.com/Microsoft/hcsshim/internal/cow"
	"github.com/Microsoft/hcsshim/internal/gcs/prot"
	hcsschema "github.com/Microsoft/hcsshim/internal/hcs/schema2"
	"github.com/Microsoft/hcsshim/internal/logfields"
	"github.com/Microsoft/hcsshim/internal/ot"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"
)

const (
	protocolVersion = 4

	firstIoChannelVsockPort = prot.LinuxGcsVsockPort + 1
	nullContainerID         = "00000000-0000-0000-0000-000000000000"
)

// IoListenFunc is a type for a function that creates a listener for a VM for
// the vsock port `port`.
type IoListenFunc func(port uint32) (net.Listener, error)

// HvsockIoListen returns an implementation of IoListenFunc that listens
// on the specified vsock port for the VM specified by `vmID`.
func HvsockIoListen(vmID guid.GUID) IoListenFunc {
	return func(port uint32) (net.Listener, error) {
		return winio.ListenHvsock(&winio.HvsockAddr{
			VMID:      vmID,
			ServiceID: winio.VsockServiceID(port),
		})
	}
}

type InitialGuestState struct {
	// Timezone is only honored for Windows guests.
	Timezone *hcsschema.TimeZoneInformation
}

// GuestConnectionConfig contains options for creating a guest connection.
type GuestConnectionConfig struct {
	// Conn specifies the connection to use for the bridge. It will be closed
	// when there is an error or Close is called.
	Conn io.ReadWriteCloser
	// Log specifies the logrus entry to use for async log messages.
	Log *logrus.Entry
	// IoListen is the function to use to create listeners for the stdio connections.
	IoListen IoListenFunc
	// InitGuestState specifies settings to apply to the guest on creation/start. This includes things such as the timezone for the VM.
	InitGuestState *InitialGuestState
}

// Connect establishes a GCS connection. `gcc.Conn` will be closed by this function.
func (gcc *GuestConnectionConfig) Connect(ctx context.Context, isColdStart bool) (_ *GuestConnection, err error) {
	ctx, span := ot.StartSpan(ctx, "gcs::GuestConnectionConfig::Connect", ot.WithClientSpanKind)
	defer span.End()
	defer func() { ot.SetSpanStatus(span, err) }()

	gc := &GuestConnection{
		nextPort:   firstIoChannelVsockPort,
		notifyChs:  make(map[string]chan struct{}),
		ioListenFn: gcc.IoListen,
	}
	gc.brdg = newBridge(gcc.Conn, gc.notify, gcc.Log)
	gc.brdg.Start()
	go func() {
		_ = gc.brdg.Wait()
		gc.clearNotifies()
	}()
	err = gc.connect(ctx, isColdStart, gcc.InitGuestState)
	if err != nil {
		gc.Close()
		return nil, err
	}
	return gc, nil
}

// GuestConnection represents a connection to the GCS.
type GuestConnection struct {
	brdg       *bridge
	ioListenFn IoListenFunc
	mu         sync.Mutex
	nextPort   uint32
	notifyChs  map[string]chan struct{}
	caps       GuestDefinedCapabilities
	os         string
}

var _ cow.ProcessHost = &GuestConnection{}

// Capabilities returns the guest's declared capabilities.
func (gc *GuestConnection) Capabilities() GuestDefinedCapabilities {
	return gc.caps
}

// Protocol returns the protocol version that is in use.
func (gc *GuestConnection) Protocol() uint32 {
	return protocolVersion
}

// connect establishes the GCS connection: it negotiates the protocol version,
// records the guest OS and capabilities, and on a cold start sends the host
// create/start messages. This method must not be called more than once per active connection.
// isColdStart should be true when the UVM is being connected to for the first time post-boot.
// It should be false for subsequent connections (e.g. if reconnecting to an existing UVM).
func (gc *GuestConnection) connect(ctx context.Context, isColdStart bool, initGuestState *InitialGuestState) (err error) {
	req := prot.NegotiateProtocolRequest{
		MinimumVersion: protocolVersion,
		MaximumVersion: protocolVersion,
	}
	var resp prot.NegotiateProtocolResponse
	err = gc.brdg.RPC(ctx, prot.RPCNegotiateProtocol, &req, &resp, true)
	if err != nil {
		return err
	}
	if resp.Version != protocolVersion {
		return fmt.Errorf("unexpected version %d returned", resp.Version)
	}

	gc.os = strings.ToLower(resp.Capabilities.RuntimeOsType)
	if gc.os == "" {
		gc.os = "windows"
	}

	gc.caps, err = unmarshalGuestCapabilities(gc.os, resp.Capabilities.GuestDefinedCapabilities)
	if err != nil {
		return fmt.Errorf("unmarshalGuestCapabilities: %w", err)
	}

	if isColdStart && resp.Capabilities.SendHostCreateMessage {
		conf := &prot.UvmConfig{
			SystemType: "Container",
		}
		if initGuestState != nil && initGuestState.Timezone != nil {
			conf.TimeZoneInformation = initGuestState.Timezone
		}
		createReq := prot.ContainerCreate{
			RequestBase:     makeRequest(ctx, nullContainerID),
			ContainerConfig: prot.AnyInString{Value: conf},
		}
		var createResp prot.ResponseBase
		err = gc.brdg.RPC(ctx, prot.RPCCreate, &createReq, &createResp, true)
		if err != nil {
			return err
		}
		if resp.Capabilities.SendHostStartMessage {
			startReq := makeRequest(ctx, nullContainerID)
			var startResp prot.ResponseBase
			err = gc.brdg.RPC(ctx, prot.RPCStart, &startReq, &startResp, true)
			if err != nil {
				return err
			}
		}
	}
	return nil
}

// Modify sends a modify settings request to the null container. This is
// generally used to prepare virtual hardware that has been added to the guest.
func (gc *GuestConnection) Modify(ctx context.Context, settings interface{}) (err error) {
	ctx, span := ot.StartSpan(ctx, "gcs::GuestConnection::Modify", ot.WithClientSpanKind)
	defer span.End()
	defer func() { ot.SetSpanStatus(span, err) }()

	req := prot.ContainerModifySettings{
		RequestBase: makeRequest(ctx, nullContainerID),
		Request:     settings,
	}
	var resp prot.ResponseBase
	return gc.brdg.RPC(ctx, prot.RPCModifySettings, &req, &resp, false)
}

func (gc *GuestConnection) ModifyServiceSettings(ctx context.Context, serviceType prot.ServiceModifyPropertyType, settings interface{}) (err error) {
	ctx, span := ot.StartSpan(ctx, "gcs::GuestConnection::ModifyServiceSettings", ot.WithClientSpanKind)
	defer span.End()
	defer func() { ot.SetSpanStatus(span, err) }()

	req := prot.ServiceModificationRequest{
		RequestBase:  makeRequest(ctx, nullContainerID),
		PropertyType: string(serviceType),
		Settings:     settings,
	}
	var resp prot.ResponseBase
	return gc.brdg.RPC(ctx, prot.RPCModifyServiceSettings, &req, &resp, false)
}

func (gc *GuestConnection) DumpStacks(ctx context.Context) (response string, err error) {
	ctx, span := ot.StartSpan(ctx, "gcs::GuestConnection::DumpStacks", ot.WithClientSpanKind)
	defer span.End()
	defer func() { ot.SetSpanStatus(span, err) }()

	req := prot.DumpStacksRequest{
		RequestBase: makeRequest(ctx, nullContainerID),
	}
	var resp prot.DumpStacksResponse
	err = gc.brdg.RPC(ctx, prot.RPCDumpStacks, &req, &resp, false)
	return resp.GuestStacks, err
}

func (gc *GuestConnection) DeleteContainerState(ctx context.Context, cid string) (err error) {
	ctx, span := ot.StartSpan(ctx, "gcs::GuestConnection::DeleteContainerState", ot.WithClientSpanKind)
	defer span.End()
	defer func() { ot.SetSpanStatus(span, err) }()
	span.SetAttributes(attribute.String("cid", cid))

	req := prot.DeleteContainerStateRequest{
		RequestBase: makeRequest(ctx, cid),
	}
	var resp prot.ResponseBase
	return gc.brdg.RPC(ctx, prot.RPCDeleteContainerState, &req, &resp, false)
}

// Close terminates the guest connection. It is undefined to call any other
// methods on the connection after this is called.
func (gc *GuestConnection) Close() error {
	if gc.brdg == nil {
		return nil
	}
	return gc.brdg.Close()
}

// SetMigrating forwards to [bridge.SetMigrating]. No-op if uninitialized.
func (gc *GuestConnection) SetMigrating(migrating bool) {
	if gc.brdg == nil {
		return
	}

	gc.brdg.SetMigrating(migrating)
}

// IsBridgeConnected reports whether a live bridge transport is currently installed.
func (gc *GuestConnection) IsBridgeConnected() bool {
	if gc.brdg == nil {
		return false
	}
	return gc.brdg.connected.Load()
}

// ResumeOnConn resumes the bridge after swaping the bridge
// transport without dropping outstanding RPCs.
func (gc *GuestConnection) ResumeOnConn(ctx context.Context, conn io.ReadWriteCloser) error {
	if gc.brdg == nil {
		// Not adopting conn; close it so the accepted socket does not leak.
		_ = conn.Close()
		return ErrBridgeClosed
	}
	if err := gc.brdg.ResumeOnConn(conn); err != nil {
		return err
	}

	// The guest resets its protocol version when it re-dials after the blackout,
	// so renegotiate before any version-gated RPC (e.g. exec) is issued.
	return gc.connect(ctx, false, nil)
}

// CreateProcess creates a process in the container host.
func (gc *GuestConnection) CreateProcess(ctx context.Context, settings interface{}) (_ cow.Process, err error) {
	ctx, span := ot.StartSpan(ctx, "gcs::GuestConnection::CreateProcess", ot.WithClientSpanKind)
	defer span.End()
	defer func() { ot.SetSpanStatus(span, err) }()

	return gc.exec(ctx, nullContainerID, settings)
}

// OS returns the operating system of the container's host, "windows" or "linux".
func (gc *GuestConnection) OS() string {
	return gc.os
}

// IsOCI returns false, indicating that CreateProcess should not be called with
// an OCI process spec.
func (gc *GuestConnection) IsOCI() bool {
	return false
}

func (gc *GuestConnection) newIoChannel() (*ioChannel, uint32, error) {
	gc.mu.Lock()
	port := gc.nextPort
	gc.nextPort++
	gc.mu.Unlock()
	l, err := gc.ioListenFn(port)
	if err != nil {
		return nil, 0, err
	}
	return newIoChannel(l), port, nil
}

// SetNextPort raises the new-process IO port allocator floor. Called
// by the live-migration restore path after [Connect] to skip past
// vsock ports already in use by restored processes. Never goes
// backwards.
func (gc *GuestConnection) SetNextPort(p uint32) {
	gc.mu.Lock()
	defer gc.mu.Unlock()
	if p > gc.nextPort {
		gc.nextPort = p
	}
}

// NextPort returns the current allocator floor. Used by the
// live-migration save path to record what [SetNextPort] should be
// seeded with on the destination.
func (gc *GuestConnection) NextPort() uint32 {
	gc.mu.Lock()
	defer gc.mu.Unlock()
	return gc.nextPort
}

// BridgeNextID returns the bridge's next request id allocator value.
func (gc *GuestConnection) BridgeNextID() int64 {
	return gc.brdg.NextID()
}

// SeedBridgeNextID raises the bridge's request id allocator floor.
func (gc *GuestConnection) SeedBridgeNextID(next int64) {
	gc.brdg.SeedNextID(next)
}

// PreregisterWaitForProcess registers a stub WaitForProcess rpc against the
// source-issued id so a later [Container.OpenProcessWithIO] (or the bridge
// itself) routes the guest's response without issuing a duplicate wait.
func (gc *GuestConnection) PreregisterWaitForProcess(id int64, resp *prot.ContainerWaitForProcessResponse) error {
	_, err := gc.brdg.PreregisterRPC(id, prot.RPCWaitForProcess, resp)
	return err
}

func (gc *GuestConnection) requestNotify(cid string, ch chan struct{}) error {
	gc.mu.Lock()
	defer gc.mu.Unlock()
	if gc.notifyChs == nil {
		return errors.New("guest connection closed")
	}
	if _, ok := gc.notifyChs[cid]; ok {
		return fmt.Errorf("container %s already exists", cid)
	}
	gc.notifyChs[cid] = ch
	return nil
}

func (gc *GuestConnection) notify(ntf *prot.ContainerNotification) error {
	cid := ntf.ContainerID
	gc.mu.Lock()
	ch := gc.notifyChs[cid]
	delete(gc.notifyChs, cid)
	gc.mu.Unlock()
	if ch == nil {
		return fmt.Errorf("container %s not found", cid)
	}
	logrus.WithField(logfields.ContainerID, cid).Info("container terminated in guest")
	close(ch)
	return nil
}

func (gc *GuestConnection) clearNotifies() {
	gc.mu.Lock()
	chs := gc.notifyChs
	gc.notifyChs = nil
	gc.mu.Unlock()
	for _, ch := range chs {
		close(ch)
	}
}

func makeRequest(ctx context.Context, cid string) prot.RequestBase {
	r := prot.RequestBase{
		ContainerID: cid,
	}
	span := trace.SpanFromContext(ctx)
	if span != nil && span.SpanContext().IsValid() {
		carrier := make(prot.Otelspancontext)
		otel.GetTextMapPropagator().Inject(ctx, propagation.MapCarrier(carrier))
		r.OpenTelemetrySpanContext = carrier
	}
	return r
}