package runtime

import (
	"context"
	"encoding/binary"
	"encoding/hex"
	"errors"
	"io"
	"net"
	"strconv"
	"strings"
	"time"
	"unicode/utf8"

	"browser.local/run/protocol"
)

const (
	sourceRCONAuthRequestID       int32 = 1
	sourceRCONCommandRequestID    int32 = 2
	sourceRCONResponseValue       int32 = 0
	sourceRCONAuthResponse        int32 = 2
	sourceRCONExecuteCommand      int32 = 2
	sourceRCONAuthenticate        int32 = 3
	sourceRCONMaxPacketSize             = 4096
	sourceRCONMaxCommandBytes           = 4000
	sourceRCONMaxResponsePackets        = 32
	sourceRCONMaxResponseBytes          = 64 * 1024
	sourceRCONIOTimeout                 = 10 * time.Second
	sourceRCONMaxExecutionTimeout       = 60 * time.Second
)

type sourceRCONPacket struct {
	id       int32
	typeCode int32
	body     string
}

type sourceRCONConfig struct {
	password string
}

type sourceRCONError struct {
	code string
}

func (err sourceRCONError) Error() string { return err.code }

// ExecuteSourceRCON connects only to the local UE4SS listener described by a
// frozen plan. The command, config, password, and response body are transient.
func (executor LifecycleExecutor) ExecuteSourceRCON(ctx context.Context, assignment protocol.RunJobAssignment, command string) LifecycleExecutionResult {
	if assignment.ExecutionInput.SourceRCON == nil || protocol.ValidateRunJobAssignment(assignment) != nil {
		return lifecycleFailure("unsafe_source_rcon_plan", "Source RCON plan is invalid")
	}
	if executor.runtimeTargetOS != "windows" || executor.runtimeTargetArch != "amd64" {
		return lifecycleFailure("unsupported_extension_platform", "Source RCON requires Windows amd64")
	}
	if !validSourceRCONCommand(command) {
		return lifecycleFailure("source_rcon_input_invalid", "Source RCON command input is invalid")
	}

	executionCtx, cancel := sourceRCONExecutionContext(ctx, assignment.ExecutionInput.TimeoutSeconds)
	defer cancel()
	resolver := NewWorkspaceResolver(executor.workspaceRoot)
	scope, err := resolver.Scope(assignment.ServerInstanceID, assignment.ExecutionInput.WorkspaceScope)
	if err != nil {
		return lifecycleFailure("source_rcon_workspace_unavailable", "Source RCON workspace is unavailable")
	}
	configPath, err := sourceRCONConfigPath(resolver, scope, *assignment.ExecutionInput.SourceRCON)
	if err != nil {
		return lifecycleFailure("source_rcon_config_unavailable", "Source RCON configuration is unavailable")
	}
	config, err := loadSourceRCONConfig(configPath, assignment.ExecutionInput.SourceRCON.Port)
	if err != nil {
		return lifecycleFailure("source_rcon_config_invalid", "Source RCON configuration is invalid")
	}
	if err := executeSourceRCONWire(executionCtx, assignment.ExecutionInput.SourceRCON.Port, config.password, command); err != nil {
		if errors.Is(executionCtx.Err(), context.Canceled) {
			return LifecycleExecutionResult{State: lifecycleResultStateCancelled, Progress: protocol.RunJobProgressReport{Percent: 100, Message: "Source RCON command cancelled"}, Message: "Source RCON command cancelled", ErrorCode: "source_rcon_cancelled"}
		}
		if errors.Is(executionCtx.Err(), context.DeadlineExceeded) {
			return lifecycleFailure("source_rcon_timeout", "Source RCON command timed out")
		}
		var sourceErr sourceRCONError
		if errors.As(err, &sourceErr) {
			return lifecycleFailure(sourceErr.code, sourceRCONSafeMessage(sourceErr.code))
		}
		return lifecycleFailure("source_rcon_execution_failed", "Source RCON command failed")
	}
	return LifecycleExecutionResult{
		State:    lifecycleResultStateSucceeded,
		Progress: protocol.RunJobProgressReport{Percent: 100, Message: "Source RCON command delivered"},
		Message:  "Source RCON command delivered",
		ExecutionResult: protocol.RunJobExecutionResult{
			Kind:    "source-rcon",
			Summary: "one-time loopback Source RCON command delivered",
		},
	}
}

func sourceRCONExecutionContext(ctx context.Context, timeoutSeconds int) (context.Context, context.CancelFunc) {
	timeout := time.Duration(timeoutSeconds) * time.Second
	if timeout <= 0 || timeout > sourceRCONMaxExecutionTimeout {
		timeout = sourceRCONMaxExecutionTimeout
	}
	return context.WithTimeout(ctx, timeout)
}

func validSourceRCONCommand(command string) bool {
	return strings.TrimSpace(command) != "" && utf8.ValidString(command) && len([]byte(command)) <= sourceRCONMaxCommandBytes && !strings.ContainsAny(command, "\x00\r\n")
}

func loadSourceRCONConfig(path string, expectedPort int) (sourceRCONConfig, error) {
	body, found, err := readBoundedRegularFile(path, maxUE4SSMetadataBytes)
	if err != nil || !found || !utf8.Valid(body) {
		return sourceRCONConfig{}, sourceRCONError{code: "source_rcon_config_invalid"}
	}
	content := strings.ReplaceAll(string(body), "\r\n", "\n")
	if !strings.Contains(content, managedRCONConfigMarker) {
		return sourceRCONConfig{}, sourceRCONError{code: "source_rcon_config_invalid"}
	}
	values := map[string]string{}
	inRCON := false
	for _, rawLine := range strings.Split(content, "\n") {
		line := strings.TrimSpace(rawLine)
		if line == "[rcon]" {
			inRCON = true
			continue
		}
		if strings.HasPrefix(line, "[") {
			inRCON = false
			continue
		}
		if !inRCON || line == "" || strings.HasPrefix(line, ";") || strings.HasPrefix(line, "#") {
			continue
		}
		key, value, ok := strings.Cut(line, "=")
		if !ok {
			continue
		}
		key = strings.TrimSpace(key)
		if key != "bind_address" && key != "port" && key != "password" {
			continue
		}
		if _, duplicate := values[key]; duplicate {
			return sourceRCONConfig{}, sourceRCONError{code: "source_rcon_config_invalid"}
		}
		values[key] = strings.TrimSpace(value)
	}
	configuredPort, err := strconv.Atoi(values["port"])
	if err != nil || values["bind_address"] != "127.0.0.1" || configuredPort != expectedPort || len(values["password"]) != 64 {
		return sourceRCONConfig{}, sourceRCONError{code: "source_rcon_config_invalid"}
	}
	if _, err := hex.DecodeString(values["password"]); err != nil {
		return sourceRCONConfig{}, sourceRCONError{code: "source_rcon_config_invalid"}
	}
	return sourceRCONConfig{password: values["password"]}, nil
}

func sourceRCONConfigPath(resolver WorkspaceResolver, scope string, plan protocol.RuntimeSourceRCONPlan) (string, error) {
	markerPath, err := resolver.ExistingTarget(scope, plan.DeploymentStateRef)
	if err != nil {
		return "", sourceRCONError{code: "source_rcon_config_unavailable"}
	}
	marker, found, err := loadManagedDLLExtensionMarker(markerPath)
	if err != nil || !found || !sourceRCONMarkerMatchesPlan(marker, plan) {
		return "", sourceRCONError{code: "source_rcon_config_unavailable"}
	}
	configPath, err := resolver.ExistingTarget(scope, marker.ConfigRef)
	if err != nil {
		return "", sourceRCONError{code: "source_rcon_config_unavailable"}
	}
	return configPath, nil
}

func sourceRCONMarkerMatchesPlan(marker managedDLLExtensionMarker, plan protocol.RuntimeSourceRCONPlan) bool {
	if marker.ExtensionKey != plan.ExtensionKey || marker.ModKey != plan.ModKey || marker.RCONPort != plan.Port || !managedRCONConfigRefForMod(marker.ConfigRef, plan.ModKey) {
		return false
	}
	return marker.ConfigRef == plan.ConfigRef || strings.HasSuffix(marker.ConfigRef, "/"+plan.ConfigRef)
}

func executeSourceRCONWire(ctx context.Context, port int, password string, command string) error {
	dialer := net.Dialer{Timeout: sourceRCONIOTimeout}
	connection, err := dialer.DialContext(ctx, "tcp4", net.JoinHostPort("127.0.0.1", strconv.Itoa(port)))
	if err != nil {
		return sourceRCONError{code: "source_rcon_connection_failed"}
	}
	defer connection.Close()
	if err := writeSourceRCONPacket(ctx, connection, sourceRCONPacket{id: sourceRCONAuthRequestID, typeCode: sourceRCONAuthenticate, body: password}); err != nil {
		return sourceRCONError{code: "source_rcon_connection_failed"}
	}
	auth, err := readSourceRCONPacket(ctx, connection)
	if err != nil {
		return sourceRCONError{code: "source_rcon_protocol_failed"}
	}
	if auth.typeCode != sourceRCONAuthResponse || auth.id == -1 {
		return sourceRCONError{code: "source_rcon_authentication_failed"}
	}
	if auth.id != sourceRCONAuthRequestID || auth.body != "" {
		return sourceRCONError{code: "source_rcon_protocol_failed"}
	}
	if err := writeSourceRCONPacket(ctx, connection, sourceRCONPacket{id: sourceRCONCommandRequestID, typeCode: sourceRCONExecuteCommand, body: command}); err != nil {
		return sourceRCONError{code: "source_rcon_connection_failed"}
	}
	responseBytes := 0
	sourceError := false
	responsePrefix := make([]byte, 0, len("error:"))
	for packetIndex := 0; packetIndex < sourceRCONMaxResponsePackets; packetIndex++ {
		response, err := readSourceRCONPacket(ctx, connection)
		if err != nil {
			return sourceRCONError{code: "source_rcon_protocol_failed"}
		}
		if response.id != sourceRCONCommandRequestID || response.typeCode != sourceRCONResponseValue {
			return sourceRCONError{code: "source_rcon_protocol_failed"}
		}
		responseBytes += len([]byte(response.body))
		if responseBytes > sourceRCONMaxResponseBytes {
			return sourceRCONError{code: "source_rcon_protocol_failed"}
		}
		if len(responsePrefix) < cap(responsePrefix) {
			remaining := cap(responsePrefix) - len(responsePrefix)
			chunk := []byte(response.body)
			if len(chunk) > remaining {
				chunk = chunk[:remaining]
			}
			responsePrefix = append(responsePrefix, chunk...)
		}
		if len(responsePrefix) == len("error:") && strings.EqualFold(string(responsePrefix), "error:") {
			sourceError = true
		}
		if response.body == "" {
			if sourceError {
				return sourceRCONError{code: "source_rcon_command_failed"}
			}
			return nil
		}
	}
	return sourceRCONError{code: "source_rcon_protocol_failed"}
}

func writeSourceRCONPacket(ctx context.Context, connection net.Conn, packet sourceRCONPacket) error {
	if !utf8.ValidString(packet.body) || len([]byte(packet.body)) > sourceRCONMaxCommandBytes {
		return sourceRCONError{code: "source_rcon_protocol_failed"}
	}
	size := 8 + len(packet.body) + 2
	if size < 10 || size > sourceRCONMaxPacketSize {
		return sourceRCONError{code: "source_rcon_protocol_failed"}
	}
	buffer := make([]byte, 4+size)
	binary.LittleEndian.PutUint32(buffer[0:4], uint32(size))
	binary.LittleEndian.PutUint32(buffer[4:8], uint32(packet.id))
	binary.LittleEndian.PutUint32(buffer[8:12], uint32(packet.typeCode))
	copy(buffer[12:], packet.body)
	if err := setSourceRCONDeadline(ctx, connection); err != nil {
		return err
	}
	_, err := connection.Write(buffer)
	return err
}

func readSourceRCONPacket(ctx context.Context, connection net.Conn) (sourceRCONPacket, error) {
	if err := setSourceRCONDeadline(ctx, connection); err != nil {
		return sourceRCONPacket{}, err
	}
	var sizeBuffer [4]byte
	if _, err := io.ReadFull(connection, sizeBuffer[:]); err != nil {
		return sourceRCONPacket{}, err
	}
	size := int(int32(binary.LittleEndian.Uint32(sizeBuffer[:])))
	if size < 10 || size > sourceRCONMaxPacketSize {
		return sourceRCONPacket{}, sourceRCONError{code: "source_rcon_protocol_failed"}
	}
	payload := make([]byte, size)
	if _, err := io.ReadFull(connection, payload); err != nil {
		return sourceRCONPacket{}, err
	}
	if payload[size-2] != 0 || payload[size-1] != 0 || !utf8.Valid(payload[8:size-2]) {
		return sourceRCONPacket{}, sourceRCONError{code: "source_rcon_protocol_failed"}
	}
	return sourceRCONPacket{id: int32(binary.LittleEndian.Uint32(payload[0:4])), typeCode: int32(binary.LittleEndian.Uint32(payload[4:8])), body: string(payload[8 : size-2])}, nil
}

func setSourceRCONDeadline(ctx context.Context, connection net.Conn) error {
	deadline := time.Now().Add(sourceRCONIOTimeout)
	if contextDeadline, ok := ctx.Deadline(); ok && contextDeadline.Before(deadline) {
		deadline = contextDeadline
	}
	return connection.SetDeadline(deadline)
}

func sourceRCONSafeMessage(code string) string {
	switch code {
	case "source_rcon_connection_failed":
		return "Source RCON listener is unavailable"
	case "source_rcon_authentication_failed":
		return "Source RCON authentication failed"
	case "source_rcon_command_failed":
		return "Source RCON command was rejected"
	case "source_rcon_protocol_failed":
		return "Source RCON protocol exchange failed"
	default:
		return "Source RCON command failed"
	}
}