package service

import (
	"context"
	"errors"
	"time"

	toki "git.toki-labs.com/toki/proto-socket/go"

	edgenode "iop/apps/edge/internal/node"
	"iop/packages/go/config"
	"iop/packages/go/workspaceprotocol"
	iop "iop/proto/gen/iop"
)

var (
	errWorkspaceWireUnavailable = errors.New("workspace wire: admitted node is unavailable")
	errWorkspaceWireStale       = errors.New("workspace wire: admitted node connection changed")
	errWorkspaceWireTransport   = errors.New("workspace wire: request failed")
	errWorkspaceWireReference   = errors.New("workspace wire: workspace reference is not admitted")
	errWorkspaceWireResponse    = errors.New("workspace wire: node response was not accepted")
	errWorkspaceWireArtifact    = errors.New("workspace wire: artifact request was not accepted")
)

// workspaceOpen sends only to the Node and connection generation frozen by
// workspace admission. It never resolves a replacement Node after reconnect, and
// it rejects any request whose workspace reference is not the exact admitted
// reference before dispatching to the Node.
func (s *Service) workspaceOpen(ctx context.Context, binding *SingleRequestWorkspaceBinding, req *iop.WorkspaceOpenRequest) (*iop.WorkspaceOpenResponse, error) {
	if req == nil || req.GetRequestId() == "" {
		return nil, errWorkspaceWireUnavailable
	}
	if binding == nil {
		return nil, errWorkspaceWireUnavailable
	}
	if req.GetWorkspaceRef() == "" || req.GetWorkspaceRef() != binding.Ref {
		return nil, errWorkspaceWireReference
	}
	outbound, err := workspaceOpenRequestFromBinding(req, binding)
	if err != nil {
		return nil, errWorkspaceWireUnavailable
	}
	wait := workspaceWireTimeout(ctx, binding, outbound.GetTimeoutMs())
	var response *iop.WorkspaceOpenResponse
	err = s.withWorkspaceBinding(binding, func(entry *edgenode.NodeEntry) error {
		var err error
		response, err = toki.SendRequestTyped[*iop.WorkspaceOpenRequest, *iop.WorkspaceOpenResponse](&entry.Client.Communicator, outbound, wait)
		return err
	})
	if err != nil {
		return nil, workspaceWireError(err)
	}
	return validateWorkspaceOpenResponse(outbound, response)
}

func workspaceOpenRequestFromBinding(req *iop.WorkspaceOpenRequest, binding *SingleRequestWorkspaceBinding) (*iop.WorkspaceOpenRequest, error) {
	frozen, err := validateAndCloneWorkspaceBinding(binding)
	if err != nil {
		return nil, err
	}
	operations := make([]iop.WorkspaceOperation, 0, len(frozen.OperationIDs))
	var readEnabled, listEnabled, writeEnabled, commandEnabled bool
	for _, operationID := range frozen.OperationIDs {
		switch config.WorkspaceOperation(operationID) {
		case config.WorkspaceOpRead:
			operations = append(operations, iop.WorkspaceOperation_WORKSPACE_OPERATION_READ)
			readEnabled = true
		case config.WorkspaceOpList:
			operations = append(operations, iop.WorkspaceOperation_WORKSPACE_OPERATION_LIST)
			listEnabled = true
		case config.WorkspaceOpWrite:
			operations = append(operations, iop.WorkspaceOperation_WORKSPACE_OPERATION_WRITE)
			writeEnabled = true
		case config.WorkspaceOpDelete:
			operations = append(operations, iop.WorkspaceOperation_WORKSPACE_OPERATION_DELETE)
		case config.WorkspaceOpCommand:
			operations = append(operations, iop.WorkspaceOperation_WORKSPACE_OPERATION_COMMAND)
			commandEnabled = true
		default:
			return nil, errSingleRequestWorkspaceMalformed
		}
	}
	outbound := &iop.WorkspaceOpenRequest{
		RequestId: req.GetRequestId(), WorkspaceRef: frozen.Ref, TimeoutMs: req.GetTimeoutMs(),
		Operations: operations, CommandIds: append([]string(nil), frozen.CommandIDs...),
	}
	if readEnabled {
		outbound.MaxReadBytes = int64(frozen.Limits.MaxReadBytes)
	}
	if writeEnabled {
		outbound.MaxWriteBytes = int64(frozen.Limits.MaxWriteBytes)
	}
	if listEnabled || commandEnabled {
		outbound.MaxOutputBytes = int64(frozen.Limits.MaxOutputBytes)
	}
	if commandEnabled {
		outbound.MaxCommandTimeoutMs = int64(frozen.Limits.MaxCommandTimeoutMS)
	}
	return outbound, nil
}

// workspaceTool uses a bounded waiter and sends one typed cancel for an
// in-flight tool call when its context is cancelled. The immutable request,
// stage, and tool-call identities are copied to that cancellation request.
//
// The admitted Node communicator is captured once up front. When the caller
// context wins the race against the tool response, cancellation is a single
// fire-and-forget typed send to that captured communicator: it never acquires
// the registry dispatch-owner mutex the in-flight tool request is holding, never
// re-selects a Node, and does not wait for a cancel response.
func (s *Service) workspaceTool(ctx context.Context, binding *SingleRequestWorkspaceBinding, req *iop.WorkspaceToolRequest) (*iop.WorkspaceToolResponse, error) {
	if req == nil || req.GetRequestId() == "" || req.GetStageId() == "" || req.GetToolCallId() == "" {
		return nil, errWorkspaceWireUnavailable
	}
	client, err := s.captureWorkspaceClient(binding)
	if err != nil {
		return nil, err
	}
	type result struct {
		response *iop.WorkspaceToolResponse
		err      error
	}
	resultCh := make(chan result, 1)
	wait := workspaceWireTimeout(ctx, binding, req.GetTimeoutMs())
	go func() {
		var response *iop.WorkspaceToolResponse
		err := s.withWorkspaceBinding(binding, func(entry *edgenode.NodeEntry) error {
			var requestErr error
			response, requestErr = toki.SendRequestTyped[*iop.WorkspaceToolRequest, *iop.WorkspaceToolResponse](&entry.Client.Communicator, req, wait)
			return requestErr
		})
		resultCh <- result{response: response, err: err}
	}()

	select {
	case result := <-resultCh:
		if result.err != nil {
			return nil, workspaceWireError(result.err)
		}
		return validateWorkspaceToolResponse(req, result.response)
	case <-ctx.Done():
		s.sendWorkspaceCancelToClient(client, binding, req)
		return nil, context.Cause(ctx)
	}
}

// workspaceArtifact dispatches the coordinator-only PLAN/REVIEW artifact
// family to the frozen Node generation. Both request and response payloads are
// bounded by the admitted output limit before they can cross their respective
// trust boundaries.
func (s *Service) workspaceArtifact(ctx context.Context, binding *SingleRequestWorkspaceBinding, req *iop.WorkspaceArtifactRequest, maxBytes int) (*iop.WorkspaceArtifactResponse, error) {
	if req == nil || req.GetRequestId() == "" || binding == nil || maxBytes < 1 || !validWorkspaceArtifactKind(req.GetKind()) || !validWorkspaceArtifactOperation(req.GetOperation()) {
		return nil, errWorkspaceWireArtifact
	}
	limit := maxBytes
	if (req.GetOperation() == iop.WorkspaceArtifactOperation_WORKSPACE_ARTIFACT_OPERATION_READ && len(req.GetContent()) != 0) || len(req.GetContent()) > limit {
		return nil, errWorkspaceWireArtifact
	}
	outbound := &iop.WorkspaceArtifactRequest{
		RequestId: req.GetRequestId(), Kind: req.GetKind(), Operation: req.GetOperation(),
		Content: append([]byte(nil), req.GetContent()...),
	}
	wait := workspaceWireTimeout(ctx, binding, 0)
	var response *iop.WorkspaceArtifactResponse
	err := s.withWorkspaceBinding(binding, func(entry *edgenode.NodeEntry) error {
		var requestErr error
		response, requestErr = toki.SendRequestTyped[*iop.WorkspaceArtifactRequest, *iop.WorkspaceArtifactResponse](&entry.Client.Communicator, outbound, wait)
		return requestErr
	})
	if err != nil {
		return nil, workspaceWireError(err)
	}
	return validateWorkspaceArtifactResponse(outbound, response, limit)
}

// sendWorkspaceCancelToClient issues exactly one fire-and-forget typed cancel to
// the captured admitted communicator, copying the immutable request/stage/tool
// identities. It runs in its own goroutine because the caller has already
// returned on context cancellation; it never resolves a replacement Node, never
// takes the registry lock, and ignores the cancel response.
func (s *Service) sendWorkspaceCancelToClient(client *toki.TcpClient, binding *SingleRequestWorkspaceBinding, req *iop.WorkspaceToolRequest) {
	if client == nil {
		return
	}
	cancel := &iop.WorkspaceCancelRequest{RequestId: req.GetRequestId(), StageId: req.GetStageId(), ToolCallId: req.GetToolCallId()}
	wait := workspaceWireTimeout(context.Background(), binding, 0)
	go func() {
		_, _ = toki.SendRequestTyped[*iop.WorkspaceCancelRequest, *iop.WorkspaceCancelResponse](&client.Communicator, cancel, wait)
	}()
}

// captureWorkspaceClient snapshots the admitted Node's communicator and confirms
// its exact connection generation before any tool dispatch. The captured client
// is the only Node a later cancellation may target.
func (s *Service) captureWorkspaceClient(binding *SingleRequestWorkspaceBinding) (*toki.TcpClient, error) {
	if binding == nil || binding.NodeID == "" || binding.ConnectionGeneration == 0 || s == nil || s.registry == nil {
		return nil, errWorkspaceWireUnavailable
	}
	entry, ok := s.registry.ReadyOwnerSnapshot(binding.NodeID)
	if !ok || entry == nil || entry.Client == nil {
		return nil, errWorkspaceWireUnavailable
	}
	if entry.ConnectionGeneration != binding.ConnectionGeneration {
		return nil, errWorkspaceWireStale
	}
	return entry.Client, nil
}

func (s *Service) workspaceCancel(ctx context.Context, binding *SingleRequestWorkspaceBinding, req *iop.WorkspaceCancelRequest) (*iop.WorkspaceCancelResponse, error) {
	if req == nil || req.GetRequestId() == "" || req.GetStageId() == "" || req.GetToolCallId() == "" {
		return nil, errWorkspaceWireUnavailable
	}
	wait := workspaceWireTimeout(ctx, binding, 0)
	var response *iop.WorkspaceCancelResponse
	err := s.withWorkspaceBinding(binding, func(entry *edgenode.NodeEntry) error {
		var err error
		response, err = toki.SendRequestTyped[*iop.WorkspaceCancelRequest, *iop.WorkspaceCancelResponse](&entry.Client.Communicator, req, wait)
		return err
	})
	if err != nil {
		return nil, workspaceWireError(err)
	}
	return validateWorkspaceCancelResponse(req, response)
}

func (s *Service) workspaceCleanup(ctx context.Context, binding *SingleRequestWorkspaceBinding, req *iop.WorkspaceCleanupRequest) (*iop.WorkspaceCleanupResponse, error) {
	if req == nil || req.GetRequestId() == "" {
		return nil, errWorkspaceWireUnavailable
	}
	wait := workspaceWireTimeout(ctx, binding, 0)
	var response *iop.WorkspaceCleanupResponse
	err := s.withWorkspaceBinding(binding, func(entry *edgenode.NodeEntry) error {
		var err error
		response, err = toki.SendRequestTyped[*iop.WorkspaceCleanupRequest, *iop.WorkspaceCleanupResponse](&entry.Client.Communicator, req, wait)
		return err
	})
	if err != nil {
		return nil, workspaceWireError(err)
	}
	return validateWorkspaceCleanupResponse(req, response)
}

// The response validators enforce the immutable coordinator identity echoes and
// the exact canonical status, error-code, and generic message triples for each
// operation using workspaceprotocol authority. Mismatched identities,
// non-canonical status/code pairs, raw text leakage, or nil responses are
// translated to errWorkspaceWireResponse. They never return, log, or
// interpolate raw Node-supplied text.
func validateWorkspaceOpenResponse(req *iop.WorkspaceOpenRequest, resp *iop.WorkspaceOpenResponse) (*iop.WorkspaceOpenResponse, error) {
	if resp == nil || resp.GetRequestId() != req.GetRequestId() || resp.GetWorkspaceRef() != req.GetWorkspaceRef() {
		return nil, errWorkspaceWireResponse
	}
	expectedErr, ok := workspaceprotocol.OpenTerminal(resp.GetStatus(), resp.GetErrorCode())
	if !ok || resp.GetError() != expectedErr {
		return nil, errWorkspaceWireResponse
	}
	return resp, nil
}

func validateWorkspaceToolResponse(req *iop.WorkspaceToolRequest, resp *iop.WorkspaceToolResponse) (*iop.WorkspaceToolResponse, error) {
	if resp == nil || resp.GetRequestId() != req.GetRequestId() || resp.GetStageId() != req.GetStageId() || resp.GetToolCallId() != req.GetToolCallId() {
		return nil, errWorkspaceWireResponse
	}
	expectedErr, ok := workspaceprotocol.ToolTerminal(resp.GetStatus(), resp.GetErrorCode())
	if !ok || resp.GetError() != expectedErr {
		return nil, errWorkspaceWireResponse
	}
	return resp, nil
}

func validateWorkspaceArtifactResponse(req *iop.WorkspaceArtifactRequest, resp *iop.WorkspaceArtifactResponse, limit int) (*iop.WorkspaceArtifactResponse, error) {
	if resp == nil || resp.GetRequestId() != req.GetRequestId() || resp.GetKind() != req.GetKind() || resp.GetOperation() != req.GetOperation() {
		return nil, errWorkspaceWireResponse
	}
	expectedErr, ok := workspaceprotocol.ArtifactTerminal(resp.GetStatus(), resp.GetErrorCode())
	if !ok || resp.GetError() != expectedErr || len(resp.GetContent()) > limit {
		return nil, errWorkspaceWireResponse
	}
	if (resp.GetStatus() != iop.WorkspaceStatus_WORKSPACE_STATUS_SUCCESS && len(resp.GetContent()) != 0) ||
		(resp.GetOperation() == iop.WorkspaceArtifactOperation_WORKSPACE_ARTIFACT_OPERATION_WRITE && len(resp.GetContent()) != 0) {
		return nil, errWorkspaceWireResponse
	}
	return resp, nil
}

func validWorkspaceArtifactKind(kind iop.WorkspaceArtifactKind) bool {
	return kind == iop.WorkspaceArtifactKind_WORKSPACE_ARTIFACT_KIND_PLAN || kind == iop.WorkspaceArtifactKind_WORKSPACE_ARTIFACT_KIND_REVIEW
}

func validWorkspaceArtifactOperation(operation iop.WorkspaceArtifactOperation) bool {
	return operation == iop.WorkspaceArtifactOperation_WORKSPACE_ARTIFACT_OPERATION_READ || operation == iop.WorkspaceArtifactOperation_WORKSPACE_ARTIFACT_OPERATION_WRITE
}

func validateWorkspaceCancelResponse(req *iop.WorkspaceCancelRequest, resp *iop.WorkspaceCancelResponse) (*iop.WorkspaceCancelResponse, error) {
	if resp == nil || resp.GetRequestId() != req.GetRequestId() || resp.GetStageId() != req.GetStageId() || resp.GetToolCallId() != req.GetToolCallId() {
		return nil, errWorkspaceWireResponse
	}
	expectedErr, ok := workspaceprotocol.CancelTerminal(resp.GetStatus(), resp.GetErrorCode())
	if !ok || resp.GetError() != expectedErr {
		return nil, errWorkspaceWireResponse
	}
	return resp, nil
}

func validateWorkspaceCleanupResponse(req *iop.WorkspaceCleanupRequest, resp *iop.WorkspaceCleanupResponse) (*iop.WorkspaceCleanupResponse, error) {
	if resp == nil || resp.GetRequestId() != req.GetRequestId() {
		return nil, errWorkspaceWireResponse
	}
	expectedErr, ok := workspaceprotocol.CleanupTerminal(resp.GetStatus(), resp.GetErrorCode())
	if !ok || resp.GetError() != expectedErr {
		return nil, errWorkspaceWireResponse
	}
	return resp, nil
}

func (s *Service) withWorkspaceBinding(binding *SingleRequestWorkspaceBinding, send func(*edgenode.NodeEntry) error) error {
	if binding == nil || binding.NodeID == "" || binding.ConnectionGeneration == 0 || s == nil || s.registry == nil {
		return errWorkspaceWireUnavailable
	}
	entry, ok := s.registry.ReadyOwnerSnapshot(binding.NodeID)
	if !ok || entry == nil || entry.Client == nil {
		return errWorkspaceWireUnavailable
	}
	if entry.ConnectionGeneration != binding.ConnectionGeneration {
		return errWorkspaceWireStale
	}
	if err := s.registry.WithCurrentDispatchOwner(binding.NodeID, entry.Client, binding.ConnectionGeneration, func() error {
		return send(entry)
	}); err != nil {
		return errWorkspaceWireStale
	}
	return nil
}

func workspaceWireTimeout(ctx context.Context, binding *SingleRequestWorkspaceBinding, requestMS int64) time.Duration {
	wait := 30 * time.Second
	if binding != nil && binding.Limits.MaxCommandTimeoutMS > 0 {
		wait = time.Duration(binding.Limits.MaxCommandTimeoutMS) * time.Millisecond
	}
	if requestMS > 0 {
		requestWait := time.Duration(requestMS) * time.Millisecond
		if requestWait < wait {
			wait = requestWait
		}
	}
	if deadline, ok := ctx.Deadline(); ok {
		if remaining := time.Until(deadline); remaining < wait {
			wait = remaining
		}
	}
	if wait <= 0 {
		return time.Millisecond
	}
	return wait
}

func workspaceWireError(err error) error {
	if errors.Is(err, errWorkspaceWireUnavailable) || errors.Is(err, errWorkspaceWireStale) {
		return err
	}
	return errWorkspaceWireTransport
}