package service

import (
	"context"
	"encoding/json"
	"errors"
	"strings"
	"sync/atomic"
	"testing"
	"time"

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

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

type scriptedInternalToolExecutor struct {
	calls         []InternalWorkspaceToolCall
	results       chan InternalWorkspaceToolResult
	seenResults   []InternalWorkspaceToolResult
	continueCount atomic.Int32
}

func newScriptedInternalToolExecutor(calls ...InternalWorkspaceToolCall) *scriptedInternalToolExecutor {
	return &scriptedInternalToolExecutor{calls: calls, results: make(chan InternalWorkspaceToolResult, len(calls)+1)}
}

func (e *scriptedInternalToolExecutor) ExecuteSingleRequest(ctx context.Context, req SingleRequestRequest, ctrl SingleRequestController) error {
	sequence := uint64(1)
	if err := ctrl.SubmitEnvelope(testEnvelope(req.RequestID, sequence, SingleRequestStatePlanning)); err != nil {
		return err
	}
	for index := range e.calls {
		call := e.calls[index].Clone()
		if call.RequestID == "" {
			call.RequestID = req.RequestID
		}
		if call.StageID == "" {
			call.StageID = "plan"
		}
		sequence++
		if err := ctrl.SubmitEnvelope(SingleRequestEnvelope{
			RequestID: req.RequestID, Sequence: sequence,
			Stage: SingleRequestStateInternalTool, SavedStage: SingleRequestStatePlanning,
			ToolCall: call,
		}); err != nil {
			return err
		}
		select {
		case result := <-e.results:
			e.seenResults = append(e.seenResults, result.Clone())
		case <-ctx.Done():
			return ctx.Err()
		}
		sequence++
		if err := ctrl.SubmitEnvelope(SingleRequestEnvelope{
			RequestID: req.RequestID, Sequence: sequence,
			Stage: SingleRequestStatePlanning, SavedStage: SingleRequestStatePlanning,
		}); err != nil {
			return err
		}
	}
	for _, stage := range []SingleRequestState{SingleRequestStateWorking, SingleRequestStateReviewing, SingleRequestStateFinalizing} {
		sequence++
		envelope := testEnvelope(req.RequestID, sequence, stage)
		if stage == SingleRequestStateFinalizing {
			envelope.Result = &SingleRequestResult{Output: "private tools completed"}
		}
		if err := ctrl.SubmitEnvelope(envelope); err != nil {
			return err
		}
	}
	return nil
}

func (e *scriptedInternalToolExecutor) ContinueInternalTool(_ context.Context, result InternalWorkspaceToolResult) error {
	e.continueCount.Add(1)
	e.results <- result.Clone()
	return nil
}

func internalLoopWorkspace() config.WorkspaceDefinition {
	return config.WorkspaceDefinition{
		Ref: "workspace-loop", Platform: "darwin", Root: "/Users/operator/project",
		Operations: []config.WorkspaceOperation{
			config.WorkspaceOpRead, config.WorkspaceOpList, config.WorkspaceOpWrite,
			config.WorkspaceOpDelete, config.WorkspaceOpCommand,
		},
		Commands:             []config.WorkspaceCommandDefinition{{ID: "test"}},
		EnvironmentAllowlist: []string{"IOP_MODE"},
		MaxReadBytes:         1024, MaxWriteBytes: 1024, MaxOutputBytes: 1024, MaxCommandTimeoutMS: 1000,
	}
}

func newInternalToolLoopService(t *testing.T, executor SingleRequestExecutor) (*Service, *toki.TcpClient) {
	t.Helper()
	edgeClient, nodeClient := workspaceWirePipe(t)
	registry := edgenode.NewRegistry()
	registry.Register(&edgenode.NodeEntry{NodeID: "node-loop", Client: edgeClient})
	service := New(registry, nil)
	service.SetNodeStore(workspaceStore("node-loop", internalLoopWorkspace()))
	service.SetSingleRequestExecutor(executor)
	return service, nodeClient
}

func TestSingleRequestInternalToolLoopRequiresOptionalContinuation(t *testing.T) {
	executor := &channelFakeExecutor{fn: func(_ context.Context, req SingleRequestRequest, ctrl SingleRequestController) error {
		if err := ctrl.SubmitEnvelope(testEnvelope(req.RequestID, 1, SingleRequestStatePlanning)); err != nil {
			return err
		}
		return ctrl.SubmitEnvelope(SingleRequestEnvelope{
			RequestID: req.RequestID, Sequence: 2,
			Stage: SingleRequestStateInternalTool, SavedStage: SingleRequestStatePlanning,
			ToolCall: &InternalWorkspaceToolCall{
				RequestID: req.RequestID, StageID: "plan", ToolCallID: "tool-1",
				Name: InternalWorkspaceToolRead, Arguments: json.RawMessage(`{"relative_path":"README.md"}`),
			},
		})
	}}
	service, node := newInternalToolLoopService(t, executor)
	var openCount atomic.Int32
	installInternalLoopOpenResponder(node, &openCount)
	handle, err := service.StartSingleRequest(context.Background(), internalLoopRequest(t, nil))
	if err != nil {
		t.Fatalf("StartSingleRequest: %v", err)
	}
	if _, err := waitForExecution(t, handle); !errors.Is(err, ErrSingleRequestInternalToolUnavailable) {
		t.Fatalf("Wait error = %v, want unavailable continuation", err)
	}
	if openCount.Load() != 0 {
		t.Fatalf("unavailable continuation opened workspace %d times", openCount.Load())
	}
}

func internalLoopRequest(t *testing.T, mutate func(*SingleRequestBinding)) SingleRequestRequest {
	t.Helper()
	binding := createTestBinding(t)
	binding.WorkspaceRef = "workspace-loop"
	if mutate != nil {
		mutate(binding)
	}
	return SingleRequestRequest{RequestID: "request-loop", Binding: binding, Prompt: "complete the private task"}
}

func installInternalLoopOpenResponder(node *toki.TcpClient, count *atomic.Int32) {
	toki.AddRequestListenerTyped[*iop.WorkspaceOpenRequest, *iop.WorkspaceOpenResponse](&node.Communicator, func(req *iop.WorkspaceOpenRequest) (*iop.WorkspaceOpenResponse, error) {
		count.Add(1)
		return &iop.WorkspaceOpenResponse{
			RequestId: req.GetRequestId(), WorkspaceRef: req.GetWorkspaceRef(),
			Status: iop.WorkspaceStatus_WORKSPACE_STATUS_SUCCESS,
		}, nil
	})
	toki.AddRequestListenerTyped[*iop.WorkspaceCleanupRequest, *iop.WorkspaceCleanupResponse](&node.Communicator, func(req *iop.WorkspaceCleanupRequest) (*iop.WorkspaceCleanupResponse, error) {
		return &iop.WorkspaceCleanupResponse{RequestId: req.GetRequestId(), Status: iop.WorkspaceStatus_WORKSPACE_STATUS_SUCCESS}, nil
	})
}

func TestSingleRequestInternalToolLoopMultipleOrdered(t *testing.T) {
	executor := newScriptedInternalToolExecutor(
		InternalWorkspaceToolCall{ToolCallID: "tool-read", Name: InternalWorkspaceToolRead, Arguments: json.RawMessage(`{"relative_path":"README.md"}`)},
		InternalWorkspaceToolCall{ToolCallID: "tool-write", Name: InternalWorkspaceToolWrite, Arguments: json.RawMessage(`{"relative_path":"result.txt","content":"done"}`)},
		InternalWorkspaceToolCall{ToolCallID: "tool-command", Name: InternalWorkspaceToolCommand, Arguments: json.RawMessage(`{"command_id":"test","environment":{"IOP_MODE":"safe"}}`)},
	)
	service, node := newInternalToolLoopService(t, executor)
	var openCount atomic.Int32
	installInternalLoopOpenResponder(node, &openCount)
	toolOrder := make(chan string, 3)
	commandTimeout := make(chan int64, 1)
	toki.AddRequestListenerTyped[*iop.WorkspaceToolRequest, *iop.WorkspaceToolResponse](&node.Communicator, func(req *iop.WorkspaceToolRequest) (*iop.WorkspaceToolResponse, error) {
		toolOrder <- req.GetToolCallId()
		response := &iop.WorkspaceToolResponse{
			RequestId: req.GetRequestId(), StageId: req.GetStageId(), ToolCallId: req.GetToolCallId(),
			Status: iop.WorkspaceStatus_WORKSPACE_STATUS_SUCCESS,
		}
		switch req.GetOperation() {
		case iop.WorkspaceOperation_WORKSPACE_OPERATION_READ:
			response.Content = []byte("source")
		case iop.WorkspaceOperation_WORKSPACE_OPERATION_COMMAND:
			commandTimeout <- req.GetTimeoutMs()
			response.Stdout = []byte("ok")
		}
		return response, nil
	})

	handle, err := service.StartSingleRequest(context.Background(), internalLoopRequest(t, nil))
	if err != nil {
		t.Fatalf("StartSingleRequest: %v", err)
	}
	waitForState(t, handle, SingleRequestStateFinalizing)
	waitForSingleRequestCleanup(t, handle)
	if err := handle.AcknowledgeTerminal(true); err != nil {
		t.Fatalf("AcknowledgeTerminal: %v", err)
	}
	result, err := waitForExecution(t, handle)
	if err != nil || result.Output != "private tools completed" {
		t.Fatalf("Wait = (%q, %v)", result.Output, err)
	}
	if openCount.Load() != 1 || executor.continueCount.Load() != 3 || len(executor.seenResults) != 3 {
		t.Fatalf("open=%d continuations=%d results=%d", openCount.Load(), executor.continueCount.Load(), len(executor.seenResults))
	}
	for index, want := range []string{"tool-read", "tool-write", "tool-command"} {
		if got := <-toolOrder; got != want {
			t.Fatalf("tool order[%d]=%q, want %q", index, got, want)
		}
		if executor.seenResults[index].ToolCallID != want || executor.seenResults[index].StageID != "plan" {
			t.Fatalf("correlated result[%d]=%+v", index, executor.seenResults[index])
		}
	}
	if got := <-commandTimeout; got <= 0 || got > int64(internalLoopWorkspace().MaxCommandTimeoutMS) {
		t.Fatalf("command timeout = %d, want within workspace maximum %d", got, internalLoopWorkspace().MaxCommandTimeoutMS)
	}
	for progress := range handle.Progress() {
		if progress.Message == InternalWorkspaceToolRead || progress.Message == InternalWorkspaceToolWrite || progress.Message == InternalWorkspaceToolCommand {
			t.Fatalf("internal tool protocol reached progress: %+v", progress)
		}
	}
}

func waitForSingleRequestCleanup(t *testing.T, handle SingleRequestExecution) {
	t.Helper()
	internal, ok := handle.(*singleRequestHandle)
	if !ok {
		t.Fatal("execution does not expose coordinator cleanup state")
	}
	deadline := time.Now().Add(2 * time.Second)
	for time.Now().Before(deadline) {
		internal.mu.Lock()
		complete := internal.cleanupComplete
		internal.mu.Unlock()
		if complete {
			return
		}
		time.Sleep(time.Millisecond)
	}
	t.Fatal("workspace cleanup did not complete")
}

func TestSingleRequestInternalToolLoopFailsClosed(t *testing.T) {
	const rawSentinel = "RAW-TOOL-SENTINEL"
	tests := []struct {
		name              string
		calls             []InternalWorkspaceToolCall
		mutateBinding     func(*SingleRequestBinding)
		respond           func(*iop.WorkspaceToolRequest) *iop.WorkspaceToolResponse
		want              error
		wantWireCalls     int32
		wantContinuations int32
	}{
		{
			name:  "identity mismatch",
			calls: []InternalWorkspaceToolCall{{RequestID: "other-request", StageID: "plan", ToolCallID: "tool-1", Name: InternalWorkspaceToolRead, Arguments: json.RawMessage(`{"relative_path":"README.md"}`)}},
			want:  ErrSingleRequestIdentityMismatch,
		},
		{
			name:  "malformed arguments",
			calls: []InternalWorkspaceToolCall{{ToolCallID: "tool-1", Name: InternalWorkspaceToolRead, Arguments: json.RawMessage(`{"relative_path":"README.md","raw":"` + rawSentinel + `"}`)}},
			want:  ErrSingleRequestInternalToolInvalidCall,
		},
		{
			name:  "capability denied",
			calls: []InternalWorkspaceToolCall{{ToolCallID: "tool-1", Name: InternalWorkspaceToolCommand, Arguments: json.RawMessage(`{"command_id":"missing"}`)}},
			want:  ErrSingleRequestInternalToolDenied,
		},
		{
			name:  "environment capability denied",
			calls: []InternalWorkspaceToolCall{{ToolCallID: "tool-1", Name: InternalWorkspaceToolCommand, Arguments: json.RawMessage(`{"command_id":"test","environment":{"NOT_ALLOWED":"value"}}`)}},
			want:  ErrSingleRequestInternalToolDenied,
		},
		{
			name: "duplicate tool id",
			calls: []InternalWorkspaceToolCall{
				{ToolCallID: "tool-1", Name: InternalWorkspaceToolRead, Arguments: json.RawMessage(`{"relative_path":"README.md"}`)},
				{ToolCallID: "tool-1", Name: InternalWorkspaceToolRead, Arguments: json.RawMessage(`{"relative_path":"README.md"}`)},
			},
			want: ErrSingleRequestInternalToolInvalidCall, wantWireCalls: 1, wantContinuations: 1,
		},
		{
			name: "iteration budget",
			calls: []InternalWorkspaceToolCall{
				{ToolCallID: "tool-1", Name: InternalWorkspaceToolRead, Arguments: json.RawMessage(`{"relative_path":"README.md"}`)},
				{ToolCallID: "tool-2", Name: InternalWorkspaceToolRead, Arguments: json.RawMessage(`{"relative_path":"README.md"}`)},
			},
			mutateBinding: func(binding *SingleRequestBinding) { binding.Limits.MaxToolIterations = 1 },
			want:          ErrSingleRequestInternalToolBudget, wantWireCalls: 1, wantContinuations: 1,
		},
		{
			name:  "stale Node response",
			calls: []InternalWorkspaceToolCall{{ToolCallID: "tool-1", Name: InternalWorkspaceToolRead, Arguments: json.RawMessage(`{"relative_path":"README.md"}`)}},
			respond: func(req *iop.WorkspaceToolRequest) *iop.WorkspaceToolResponse {
				return &iop.WorkspaceToolResponse{RequestId: req.GetRequestId(), StageId: req.GetStageId(), ToolCallId: "stale-tool", Status: iop.WorkspaceStatus_WORKSPACE_STATUS_SUCCESS}
			},
			want: ErrSingleRequestInternalToolFailed, wantWireCalls: 1,
		},
		{
			name:          "output budget",
			calls:         []InternalWorkspaceToolCall{{ToolCallID: "tool-1", Name: InternalWorkspaceToolRead, Arguments: json.RawMessage(`{"relative_path":"README.md"}`)}},
			mutateBinding: func(binding *SingleRequestBinding) { binding.Limits.MaxOutputBytes = 3 },
			respond: func(req *iop.WorkspaceToolRequest) *iop.WorkspaceToolResponse {
				return &iop.WorkspaceToolResponse{RequestId: req.GetRequestId(), StageId: req.GetStageId(), ToolCallId: req.GetToolCallId(), Status: iop.WorkspaceStatus_WORKSPACE_STATUS_SUCCESS, Content: []byte("four")}
			},
			want: ErrSingleRequestInternalToolBudget, wantWireCalls: 1,
		},
		{
			name: "cumulative output budget",
			calls: []InternalWorkspaceToolCall{
				{ToolCallID: "tool-1", Name: InternalWorkspaceToolRead, Arguments: json.RawMessage(`{"relative_path":"README.md"}`)},
				{ToolCallID: "tool-2", Name: InternalWorkspaceToolRead, Arguments: json.RawMessage(`{"relative_path":"README.md"}`)},
			},
			mutateBinding: func(binding *SingleRequestBinding) { binding.Limits.MaxOutputBytes = 3 },
			respond: func(req *iop.WorkspaceToolRequest) *iop.WorkspaceToolResponse {
				return &iop.WorkspaceToolResponse{RequestId: req.GetRequestId(), StageId: req.GetStageId(), ToolCallId: req.GetToolCallId(), Status: iop.WorkspaceStatus_WORKSPACE_STATUS_SUCCESS, Content: []byte("xx")}
			},
			want: ErrSingleRequestInternalToolBudget, wantWireCalls: 2, wantContinuations: 1,
		},
	}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			executor := newScriptedInternalToolExecutor(test.calls...)
			service, node := newInternalToolLoopService(t, executor)
			var openCount atomic.Int32
			installInternalLoopOpenResponder(node, &openCount)
			var toolCount atomic.Int32
			toki.AddRequestListenerTyped[*iop.WorkspaceToolRequest, *iop.WorkspaceToolResponse](&node.Communicator, func(req *iop.WorkspaceToolRequest) (*iop.WorkspaceToolResponse, error) {
				toolCount.Add(1)
				if test.respond != nil {
					return test.respond(req), nil
				}
				return &iop.WorkspaceToolResponse{RequestId: req.GetRequestId(), StageId: req.GetStageId(), ToolCallId: req.GetToolCallId(), Status: iop.WorkspaceStatus_WORKSPACE_STATUS_SUCCESS}, nil
			})
			handle, err := service.StartSingleRequest(context.Background(), internalLoopRequest(t, test.mutateBinding))
			if err != nil {
				t.Fatalf("StartSingleRequest: %v", err)
			}
			_, err = waitForExecution(t, handle)
			if !errors.Is(err, test.want) {
				t.Fatalf("Wait error = %v, want %v", err, test.want)
			}
			if strings.Contains(err.Error(), rawSentinel) {
				t.Fatalf("raw call leaked in error %q", err)
			}
			wantOpenCalls := int32(0)
			if test.wantWireCalls > 0 {
				wantOpenCalls = 1
			}
			if openCount.Load() != wantOpenCalls || toolCount.Load() != test.wantWireCalls || executor.continueCount.Load() != test.wantContinuations {
				t.Fatalf("open=%d wire calls=%d continuations=%d, want %d/%d/%d", openCount.Load(), toolCount.Load(), executor.continueCount.Load(), wantOpenCalls, test.wantWireCalls, test.wantContinuations)
			}
		})
	}
}

func TestSingleRequestInternalToolLoopCancelPropagates(t *testing.T) {
	executor := newScriptedInternalToolExecutor(InternalWorkspaceToolCall{
		ToolCallID: "tool-command", Name: InternalWorkspaceToolCommand,
		Arguments: json.RawMessage(`{"command_id":"test"}`),
	})
	service, node := newInternalToolLoopService(t, executor)
	var openCount atomic.Int32
	installInternalLoopOpenResponder(node, &openCount)
	var sequence atomic.Int32
	toolEntered := make(chan struct{})
	release := make(chan struct{})
	cancelReached := make(chan *iop.WorkspaceCancelRequest, 1)
	serveWorkspaceConcurrent(&node.Communicator, &sequence, func(req *iop.WorkspaceToolRequest) *iop.WorkspaceToolResponse {
		close(toolEntered)
		<-release
		return &iop.WorkspaceToolResponse{RequestId: req.GetRequestId(), StageId: req.GetStageId(), ToolCallId: req.GetToolCallId(), Status: iop.WorkspaceStatus_WORKSPACE_STATUS_CANCELLED, ErrorCode: iop.WorkspaceErrorCode_WORKSPACE_ERROR_CODE_CANCELLED, Error: "workspace command cancelled"}
	})
	serveWorkspaceConcurrent(&node.Communicator, &sequence, func(req *iop.WorkspaceCancelRequest) *iop.WorkspaceCancelResponse {
		cancelReached <- req
		return &iop.WorkspaceCancelResponse{RequestId: req.GetRequestId(), StageId: req.GetStageId(), ToolCallId: req.GetToolCallId(), Status: iop.WorkspaceStatus_WORKSPACE_STATUS_CANCELLED, ErrorCode: iop.WorkspaceErrorCode_WORKSPACE_ERROR_CODE_CANCELLED, Error: "workspace command cancelled"}
	})

	handle, err := service.StartSingleRequest(context.Background(), internalLoopRequest(t, nil))
	if err != nil {
		t.Fatalf("StartSingleRequest: %v", err)
	}
	select {
	case <-toolEntered:
	case <-time.After(2 * time.Second):
		close(release)
		t.Fatal("tool did not reach Node")
	}
	handle.Cancel()
	select {
	case cancel := <-cancelReached:
		if cancel.GetRequestId() != "request-loop" || cancel.GetStageId() != "plan" || cancel.GetToolCallId() != "tool-command" {
			close(release)
			t.Fatalf("cancel identity = %+v", cancel)
		}
	case <-time.After(2 * time.Second):
		close(release)
		t.Fatal("typed cancel did not reach Node")
	}
	close(release)
	if _, err := waitForExecution(t, handle); !errors.Is(err, ErrSingleRequestCancelled) {
		t.Fatalf("Wait error = %v, want cancelled", err)
	}
	if executor.continueCount.Load() != 0 {
		t.Fatalf("cancelled call delivered %d continuations", executor.continueCount.Load())
	}
}

func TestSingleRequestInternalToolLoopStageDeadline(t *testing.T) {
	executor := newScriptedInternalToolExecutor(InternalWorkspaceToolCall{
		ToolCallID: "tool-command", Name: InternalWorkspaceToolCommand,
		Arguments: json.RawMessage(`{"command_id":"test"}`),
	})
	service, node := newInternalToolLoopService(t, executor)
	var openCount atomic.Int32
	installInternalLoopOpenResponder(node, &openCount)
	var sequence atomic.Int32
	toolEntered := make(chan struct{})
	release := make(chan struct{})
	serveWorkspaceConcurrent(&node.Communicator, &sequence, func(req *iop.WorkspaceToolRequest) *iop.WorkspaceToolResponse {
		close(toolEntered)
		<-release
		return &iop.WorkspaceToolResponse{RequestId: req.GetRequestId(), StageId: req.GetStageId(), ToolCallId: req.GetToolCallId(), Status: iop.WorkspaceStatus_WORKSPACE_STATUS_CANCELLED, ErrorCode: iop.WorkspaceErrorCode_WORKSPACE_ERROR_CODE_CANCELLED, Error: "workspace command cancelled"}
	})

	handle, err := service.StartSingleRequest(context.Background(), internalLoopRequest(t, func(binding *SingleRequestBinding) {
		binding.Limits.StageTimeoutMS = 50
	}))
	if err != nil {
		t.Fatalf("StartSingleRequest: %v", err)
	}
	select {
	case <-toolEntered:
	case <-time.After(2 * time.Second):
		close(release)
		t.Fatal("tool did not reach Node")
	}
	if _, err := waitForExecution(t, handle); !errors.Is(err, ErrSingleRequestInternalToolBudget) {
		close(release)
		t.Fatalf("Wait error = %v, want budget exhaustion", err)
	}
	close(release)
	if handle.State() != SingleRequestStateFailed || executor.continueCount.Load() != 0 {
		t.Fatalf("state=%s continuations=%d, want failed/0", handle.State(), executor.continueCount.Load())
	}
}