package wire

import (
	"context"
	"encoding/json"
	"errors"
	"log/slog"
	"strings"
	"time"
)

func (c *Protocol) handleRemoteRequestAssignment(ctx context.Context, link *elasticRequestLink, raw []byte) {
	var assignment remoteRequestAssignmentMsg
	if err := json.Unmarshal(raw, &assignment); err != nil {
		slog.Debug("AgentTeam: remote request assignment parse error", "channel", c.id, "error", err)
		return
	}
	requestID := remoteRequestID(assignment)
	if requestID == "" {
		slog.Debug("AgentTeam: remote request assignment missing requestId", "channel", c.id)
		return
	}
	if link != nil && !link.tryStart(requestID) {
		c.sendRemoteRequestResult(link, remoteRequestResultMsg{Type: remoteRequestResultType, RequestID: requestID, RemoteRequestID: requestID, Success: false, Outcome: "desktop_busy", Error: "desktop_busy"})
		return
	}
	started := c.startConnectionRunTask(ctx, func(ctx context.Context) {
		if link != nil {
			defer link.finish()
		}
		result := c.executeRemoteRequestAssignment(ctx, link, assignment)
		if link != nil && remoteRequestOutcomeClosesLink(result.Outcome) {
			link.requestDrain()
		}
		c.sendRemoteRequestResult(link, result)
	})
	if !started && link != nil {
		// tryStart reserved the link before task admission. A connection-run close
		// that wins the publication race must release that reservation.
		link.finish()
	}
}

func remoteRequestOutcomeClosesLink(outcome string) bool {
	switch strings.TrimSpace(outcome) {
	case "canceled", "deadline_exceeded", "protocol_error":
		return true
	default:
		return false
	}
}

func (c *Protocol) executeRemoteRequestAssignment(ctx context.Context, link *elasticRequestLink, assignment remoteRequestAssignmentMsg) remoteRequestResultMsg {
	requestID := remoteRequestID(assignment)
	result := remoteRequestResultMsg{Type: remoteRequestResultType, SchemaVersion: 1, RequestID: requestID, RemoteRequestID: requestID}
	if err := validateRemoteRequestAssignmentEnvelope(assignment); err != nil {
		result.Outcome = "protocol_error"
		result.Error = "protocol_error"
		return result
	}
	// Per-request endpoint authorization happens on the assignment envelope. The
	// link token only established a carrier and cannot expand this allowlist.
	if link != nil && !link.allowsEndpoint(assignment.EndpointID, assignment.EndpointVersion) {
		result.Outcome = "protocol_error"
		result.Error = "protocol_error"
		return result
	}
	deadline, _ := parseRemoteRequestDeadline(assignment)
	// Clamp execution to the earlier of the absolute deadline and relative timeout
	// budget so clock skew or stale assignments cannot extend Desktop work.
	budgetDeadline := time.Now().Add(time.Duration(assignment.TimeoutBudgetMs) * time.Millisecond)
	if budgetDeadline.Before(deadline) {
		deadline = budgetDeadline
	}
	requestCtx, cancel := context.WithDeadline(ctx, deadline)
	defer cancel()
	if link != nil {
		link.setCurrentCancel(assignment.CancelToken, cancel)
		defer link.setCurrentCancel("", nil)
	}

	registry := c.remoteEndpointRegistry
	if registry == nil {
		registry = defaultDesktopRemoteEndpointRegistry()
	}
	descriptor, ok := registry.lookup(assignment.EndpointID, assignment.EndpointVersion)
	if !ok {
		result.Outcome = "protocol_error"
		result.Error = "protocol_error"
		return result
	}
	payload, outcome, err := descriptor.Handler(requestCtx, c, assignment)
	if err != nil {
		if errors.Is(requestCtx.Err(), context.DeadlineExceeded) {
			result.Outcome = "deadline_exceeded"
			result.Error = "deadline_exceeded"
			return result
		}
		if errors.Is(requestCtx.Err(), context.Canceled) {
			result.Outcome = "canceled"
			result.Error = "canceled"
			return result
		}
		if outcome == "" {
			outcome = "endpoint_error"
		}
		result.Outcome = outcome
		result.Error = outcome
		return result
	}
	if assignment.Limits.ResponseBytes > 0 && int64(len(payload)) > assignment.Limits.ResponseBytes {
		result.Outcome = "protocol_error"
		result.Error = "protocol_error"
		return result
	}
	result.Success = true
	result.Outcome = "ok"
	result.Payload = payload
	return result
}

func validateRemoteRequestAssignmentEnvelope(assignment remoteRequestAssignmentMsg) error {
	if assignment.Type != remoteRequestAssignType || assignment.SchemaVersion != 1 || remoteRequestID(assignment) == "" || strings.TrimSpace(assignment.EndpointID) == "" || assignment.EndpointVersion <= 0 || assignment.PayloadSchemaVersion <= 0 || strings.TrimSpace(assignment.CancelToken) == "" {
		return errors.New("invalid_remote_request_envelope")
	}
	if _, err := parseRemoteRequestDeadline(assignment); err != nil {
		return err
	}
	if assignment.TimeoutBudgetMs <= 0 || assignment.Limits.RequestBytes <= 0 || assignment.Limits.ResponseBytes <= 0 {
		return errors.New("missing_remote_request_budget")
	}
	if int64(len(assignment.Payload)) > assignment.Limits.RequestBytes {
		return errors.New("remote_request_payload_too_large")
	}
	return nil
}

func parseRemoteRequestDeadline(assignment remoteRequestAssignmentMsg) (time.Time, error) {
	value := strings.TrimSpace(assignment.DeadlineAt)
	if value == "" {
		return time.Time{}, errors.New("remote_request_deadline_required")
	}
	deadline, err := time.Parse(time.RFC3339, value)
	if err != nil {
		return time.Time{}, err
	}
	return deadline.UTC(), nil
}

func remoteRequestID(assignment remoteRequestAssignmentMsg) string {
	if id := strings.TrimSpace(assignment.RequestID); id != "" {
		return id
	}
	return strings.TrimSpace(assignment.RemoteRequestID)
}

func (c *Protocol) sendRemoteRequestResult(link *elasticRequestLink, result remoteRequestResultMsg) {
	if result.Type == "" {
		result.Type = remoteRequestResultType
	}
	if result.RemoteRequestID == "" {
		result.RemoteRequestID = result.RequestID
	}
	if result.SchemaVersion <= 0 {
		result.SchemaVersion = 1
	}
	if link != nil {
		if err := link.sendJSON(result); err != nil {
			slog.Debug("AgentTeam: failed to send Elastic Request Link result", "channel", c.id, "linkId", link.linkID, "requestId", result.RequestID, "reason", "send_failed")
		}
		return
	}
	if err := c.sendJSONCritical(result); err != nil {
		slog.Debug("AgentTeam: failed to send Session Link remote request result", "channel", c.id, "requestId", result.RequestID, "reason", "send_failed")
	}
}