package daemon

import (
	"context"
	"encoding/json"
	"errors"
	"github.com/tidbcloud/AgentTeam/internal/observability/logging"
	"strings"
	"time"
)

// RemoteRequestAssignmentSpec is the concrete per-request authorization and
// execution budget sent over a Session Link Remote Request or an Elastic Request
// Link. The link grant only authorizes a carrier; this envelope authorizes one
// endpoint invocation and supplies its deadline, cancel token, and byte limits.
type RemoteRequestAssignmentSpec struct {
	RemoteRequestID      string
	EndpointID           string
	EndpointVersion      int
	PayloadSchemaVersion int
	Payload              json.RawMessage
	DeadlineAt           time.Time
	TimeoutBudget        time.Duration
	CancelToken          string
	Limits               RemoteEndpointLimits
}

type elasticRequestAssignmentState struct {
	requestID   string
	linkID      string
	endpointID  string
	deadlineAt  time.Time
	cancelToken string
	resultCh    chan RemoteRequestResult[json.RawMessage]
	done        chan struct{}
}

// AssignRemoteRequest enforces V1 link-scoped single-flight. One Elastic Request
// Link carries at most one assignment at a time; parallelism comes from multiple
// links, not multiplexing on a single carrier.
func (manager *ElasticRequestLinkManager) AssignRemoteRequest(ctx context.Context, linkID string, spec RemoteRequestAssignmentSpec) (<-chan RemoteRequestResult[json.RawMessage], error) {
	if manager == nil {
		return nil, ElasticRequestLinkError{Reason: ElasticRequestLinkRejectUnknownLink}
	}
	manager.mu.Lock()
	link := manager.links[strings.TrimSpace(linkID)]
	if link == nil {
		manager.mu.Unlock()
		return nil, ElasticRequestLinkError{Reason: ElasticRequestLinkRejectUnknownLink}
	}
	if link.currentAssignment != "" {
		manager.mu.Unlock()
		return nil, ElasticRequestLinkError{Reason: ElasticRequestLinkRejectCapacity}
	}
	spec.RemoteRequestID = strings.TrimSpace(spec.RemoteRequestID)
	spec.EndpointID = strings.TrimSpace(spec.EndpointID)
	spec.CancelToken = strings.TrimSpace(spec.CancelToken)
	if spec.PayloadSchemaVersion <= 0 {
		spec.PayloadSchemaVersion = spec.EndpointVersion
	}
	if spec.RemoteRequestID == "" || spec.EndpointID == "" || spec.EndpointVersion <= 0 || spec.PayloadSchemaVersion <= 0 || spec.DeadlineAt.IsZero() || spec.TimeoutBudget <= 0 || spec.CancelToken == "" || spec.Limits.RequestBytes <= 0 || spec.Limits.ResponseBytes <= 0 {
		manager.mu.Unlock()
		return nil, errors.New("remote request assignment requires explicit id, endpoint, deadline, budget, cancel token, and limits")
	}
	if !grantAllowsEndpoint(link.grant, spec.EndpointID, spec.EndpointVersion) {
		manager.mu.Unlock()
		return nil, ElasticRequestLinkError{Reason: ElasticRequestLinkRejectRegistryMismatch}
	}
	resultCh := make(chan RemoteRequestResult[json.RawMessage], 1)
	assignment := &elasticRequestAssignmentState{requestID: spec.RemoteRequestID, linkID: link.grant.LinkID, endpointID: spec.EndpointID, deadlineAt: spec.DeadlineAt.UTC(), cancelToken: spec.CancelToken, resultCh: resultCh, done: make(chan struct{})}
	link.currentAssignment = spec.RemoteRequestID
	manager.assignmentsByID[spec.RemoteRequestID] = assignment
	sender := link.sender
	envelope := NewRemoteRequestAssignment(spec)
	manager.recordElasticEvent("assignment_created", "ok")
	manager.recordPoolLocked()
	manager.mu.Unlock()

	if sender != nil {
		if err := sender.SendJSON(envelope); err != nil {
			manager.CleanupAssignment(spec.RemoteRequestID, "send_failed")
			return nil, err
		}
	}
	if ctx != nil {
		requestID := spec.RemoteRequestID
		done := assignment.done
		logging.Go(nil, func() {
			select {
			case <-ctx.Done():
				manager.CleanupAssignment(requestID, "canceled")
			case <-done:
			}
		})
	}
	return resultCh, nil
}

func (manager *ElasticRequestLinkManager) WaitForLink(ctx context.Context, linkID string) error {
	if manager == nil {
		return ElasticRequestLinkError{Reason: ElasticRequestLinkRejectUnknownLink}
	}
	linkID = strings.TrimSpace(linkID)
	if linkID == "" {
		return ElasticRequestLinkError{Reason: ElasticRequestLinkRejectUnknownLink}
	}
	ticker := time.NewTicker(10 * time.Millisecond)
	defer ticker.Stop()
	for {
		manager.mu.Lock()
		link := manager.links[linkID]
		manager.mu.Unlock()
		if link != nil {
			return nil
		}
		select {
		case <-ctx.Done():
			manager.CleanupExpired(manager.config.Now())
			return ctx.Err()
		case <-ticker.C:
		}
	}
}

func (manager *ElasticRequestLinkManager) CleanupLinkGrant(linkID string, reasons ...string) bool {
	if manager == nil {
		return false
	}
	reason := "manual"
	if len(reasons) > 0 && strings.TrimSpace(reasons[0]) != "" {
		reason = reasons[0]
	}
	linkID = strings.TrimSpace(linkID)
	if linkID == "" {
		return false
	}
	manager.mu.Lock()
	defer manager.mu.Unlock()
	grantState := manager.grants[linkID]
	if grantState == nil || grantState.consumed {
		return false
	}
	if warmFailureCounts(reason) {
		manager.recordWarmFailureLocked(grantState.grant, reason, manager.config.Now().UTC())
	}
	delete(manager.grants, linkID)
	manager.recordElasticEvent("grant_cleaned", reason)
	manager.recordPoolLocked()
	return true
}

func NewRemoteRequestAssignment(spec RemoteRequestAssignmentSpec) RemoteRequestAssignment[json.RawMessage] {
	deadlineAt := spec.DeadlineAt.UTC()
	deadline := deadlineAt.Format(time.RFC3339Nano)
	return RemoteRequestAssignment[json.RawMessage]{
		Type:                 sessionLinkRemoteRequestAssignType,
		SchemaVersion:        1,
		RequestID:            spec.RemoteRequestID,
		RemoteRequestID:      spec.RemoteRequestID,
		EndpointID:           spec.EndpointID,
		EndpointVersion:      spec.EndpointVersion,
		PayloadSchemaVersion: spec.PayloadSchemaVersion,
		DeadlineAt:           deadline,
		TimeoutBudgetMs:      spec.TimeoutBudget.Milliseconds(),
		CancelToken:          spec.CancelToken,
		Limits:               spec.Limits,
		Payload:              spec.Payload,
	}
}

func (manager *ElasticRequestLinkManager) CompleteRemoteRequest(result RemoteRequestResult[json.RawMessage]) bool {
	if manager == nil {
		return false
	}
	requestID := strings.TrimSpace(result.RequestID)
	if requestID == "" {
		requestID = strings.TrimSpace(result.RemoteRequestID)
	}
	manager.mu.Lock()
	assignment := manager.assignmentsByID[requestID]
	if assignment == nil {
		manager.mu.Unlock()
		return false
	}
	delete(manager.assignmentsByID, requestID)
	if link := manager.links[assignment.linkID]; link != nil {
		link.currentAssignment = ""
		link.lastIdleAt = manager.config.Now().UTC()
		if link.draining {
			sender := link.sender
			manager.removeLinkLocked(link.grant.LinkID)
			manager.recordElasticEvent("cleanup_closed_link", "drain_after_assignment")
			if link.grant.Purpose == ElasticLinkPurposeWarm {
				manager.recordWarmEvent("user_interactive", link.grant.ProfileName, "closed", "drain_after_assignment")
			}
			manager.recordPoolLocked()
			manager.mu.Unlock()
			closeAssignmentDone(assignment)
			select {
			case assignment.resultCh <- result:
			default:
			}
			if sender != nil {
				_ = sender.Close()
			}
			return true
		}
	}
	manager.recordElasticEvent("assignment_completed", "completed")
	manager.recordPoolLocked()
	manager.mu.Unlock()
	closeAssignmentDone(assignment)
	select {
	case assignment.resultCh <- result:
	default:
	}
	return true
}

func (manager *ElasticRequestLinkManager) CleanupAssignment(requestID string, reason string) bool {
	if manager == nil {
		return false
	}
	manager.mu.Lock()
	assignment := manager.assignmentsByID[strings.TrimSpace(requestID)]
	if assignment == nil {
		manager.mu.Unlock()
		return false
	}
	delete(manager.assignmentsByID, assignment.requestID)
	if link := manager.links[assignment.linkID]; link != nil {
		link.currentAssignment = ""
		link.lastIdleAt = manager.config.Now().UTC()
	}
	manager.recordElasticEvent("assignment_cleaned", reason)
	manager.recordPoolLocked()
	manager.mu.Unlock()
	closeAssignmentDone(assignment)
	result := RemoteRequestResult[json.RawMessage]{Type: sessionLinkRemoteRequestResultType, RequestID: assignment.requestID, RemoteRequestID: assignment.requestID, Success: false, Outcome: stableOutcomeForCleanup(reason), Error: stableOutcomeForCleanup(reason)}
	select {
	case assignment.resultCh <- result:
	default:
	}
	return true
}

func closeAssignmentDone(assignment *elasticRequestAssignmentState) {
	if assignment == nil || assignment.done == nil {
		return
	}
	select {
	case <-assignment.done:
	default:
		close(assignment.done)
	}
}

func stableOutcomeForCleanup(reason string) string {
	switch strings.TrimSpace(reason) {
	case "deadline_exceeded", "assignment_timeout", "connect_timeout", "idle_timeout":
		return "deadline_exceeded"
	case "canceled", "session_generation_changed", "session_disconnected":
		return "canceled"
	default:
		return "desktop_unavailable"
	}
}