package service

import (
	"context"
	"errors"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	iop "iop/proto/gen/iop"
)

type artifactLifecycleRuntime struct {
	openCount     atomic.Int32
	toolCount     atomic.Int32
	artifactCount atomic.Int32
	cleanupCount  atomic.Int32
	artifactStart chan struct{}
	artifactGate  chan struct{}
	artifactOnce  sync.Once
	mu            sync.Mutex
	artifacts     map[iop.WorkspaceArtifactKind][]byte
}

func newArtifactLifecycleRuntime(block bool) *artifactLifecycleRuntime {
	runtime := &artifactLifecycleRuntime{artifactStart: make(chan struct{}), artifacts: make(map[iop.WorkspaceArtifactKind][]byte)}
	if block {
		runtime.artifactGate = make(chan struct{})
	}
	return runtime
}

func (r *artifactLifecycleRuntime) workspaceOpen(_ context.Context, _ *SingleRequestWorkspaceBinding, req *iop.WorkspaceOpenRequest) (*iop.WorkspaceOpenResponse, error) {
	r.openCount.Add(1)
	return &iop.WorkspaceOpenResponse{RequestId: req.GetRequestId(), WorkspaceRef: req.GetWorkspaceRef(), Status: iop.WorkspaceStatus_WORKSPACE_STATUS_SUCCESS}, nil
}

func (r *artifactLifecycleRuntime) workspaceTool(_ context.Context, _ *SingleRequestWorkspaceBinding, req *iop.WorkspaceToolRequest) (*iop.WorkspaceToolResponse, error) {
	r.toolCount.Add(1)
	return &iop.WorkspaceToolResponse{RequestId: req.GetRequestId(), StageId: req.GetStageId(), ToolCallId: req.GetToolCallId(), Status: iop.WorkspaceStatus_WORKSPACE_STATUS_SUCCESS, Content: []byte("tool result")}, nil
}

func (r *artifactLifecycleRuntime) workspaceArtifact(_ context.Context, _ *SingleRequestWorkspaceBinding, req *iop.WorkspaceArtifactRequest, maximum int) (*iop.WorkspaceArtifactResponse, error) {
	r.artifactCount.Add(1)
	r.artifactOnce.Do(func() { close(r.artifactStart) })
	if r.artifactGate != nil {
		<-r.artifactGate
	}
	response := &iop.WorkspaceArtifactResponse{RequestId: req.GetRequestId(), Kind: req.GetKind(), Operation: req.GetOperation(), Status: iop.WorkspaceStatus_WORKSPACE_STATUS_SUCCESS}
	r.mu.Lock()
	defer r.mu.Unlock()
	switch req.GetOperation() {
	case iop.WorkspaceArtifactOperation_WORKSPACE_ARTIFACT_OPERATION_WRITE:
		if len(req.GetContent()) > maximum {
			return nil, errWorkspaceWireArtifact
		}
		r.artifacts[req.GetKind()] = append([]byte(nil), req.GetContent()...)
	case iop.WorkspaceArtifactOperation_WORKSPACE_ARTIFACT_OPERATION_READ:
		response.Content = append([]byte(nil), r.artifacts[req.GetKind()]...)
	}
	return response, nil
}

func (r *artifactLifecycleRuntime) CleanupWorkspace(context.Context, *SingleRequestWorkspaceBinding, string) error {
	r.cleanupCount.Add(1)
	return nil
}

type artifactAndToolExecutor struct {
	results chan InternalWorkspaceToolResult
}

func (e *artifactAndToolExecutor) ExecuteSingleRequest(ctx context.Context, req SingleRequestRequest, ctrl SingleRequestController) error {
	if err := ctrl.SubmitEnvelope(testEnvelope(req.RequestID, 1, SingleRequestStatePlanning)); err != nil {
		return err
	}
	if err := ctrl.WriteInternalArtifact(ctx, SingleRequestArtifactPlan, []byte("small plan")); err != nil {
		return err
	}
	content, err := ctrl.ReadInternalArtifact(ctx, SingleRequestArtifactPlan)
	if err != nil || string(content) != "small plan" {
		return ErrSingleRequestInternalArtifactFailed
	}
	if err := ctrl.SubmitEnvelope(SingleRequestEnvelope{
		RequestID: req.RequestID, Sequence: 2, Stage: SingleRequestStateInternalTool, SavedStage: SingleRequestStatePlanning,
		ToolCall: &InternalWorkspaceToolCall{RequestID: req.RequestID, StageID: "plan", ToolCallID: "tool-after-artifact", Name: InternalWorkspaceToolRead, Arguments: []byte(`{"relative_path":"README.md"}`)},
	}); err != nil {
		return err
	}
	select {
	case <-e.results:
	case <-ctx.Done():
		return ctx.Err()
	}
	if err := ctrl.SubmitEnvelope(SingleRequestEnvelope{RequestID: req.RequestID, Sequence: 3, Stage: SingleRequestStatePlanning, SavedStage: SingleRequestStatePlanning}); err != nil {
		return err
	}
	for sequence, stage := range []SingleRequestState{SingleRequestStateWorking, SingleRequestStateReviewing, SingleRequestStateFinalizing} {
		envelope := testEnvelope(req.RequestID, uint64(sequence+4), stage)
		if stage == SingleRequestStateFinalizing {
			envelope.Result = &SingleRequestResult{Output: "artifact lifecycle complete"}
		}
		if err := ctrl.SubmitEnvelope(envelope); err != nil {
			return err
		}
	}
	return nil
}

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

func artifactLifecycleRequest(t *testing.T, requestID string) SingleRequestRequest {
	t.Helper()
	binding := createTestBinding(t)
	binding.Workspace = &SingleRequestWorkspaceBinding{
		Ref: "workspace-ref-123", NodeID: "node-artifact", ConnectionGeneration: 11,
		OperationIDs: []string{"read"},
		Limits:       SingleRequestWorkspaceLimits{MaxReadBytes: 1024, MaxOutputBytes: 64},
	}
	return SingleRequestRequest{RequestID: requestID, Binding: binding, Prompt: "exercise artifacts"}
}

func TestSingleRequestArtifactLifecycle(t *testing.T) {
	t.Run("artifact first and tool after artifact share one open", func(t *testing.T) {
		runtime := newArtifactLifecycleRuntime(false)
		executor := &artifactAndToolExecutor{results: make(chan InternalWorkspaceToolResult, 1)}
		handle, err := startSingleRequestWithToolLoop(context.Background(), executor, executor, runtime, artifactLifecycleRequest(t, "request-artifact"))
		if err != nil {
			t.Fatal(err)
		}
		waitForState(t, handle, SingleRequestStateFinalizing)
		waitForSingleRequestCleanup(t, handle)
		if err := handle.AcknowledgeTerminal(true); err != nil {
			t.Fatal(err)
		}
		result, err := waitForExecution(t, handle)
		if err != nil || result.Output != "artifact lifecycle complete" {
			t.Fatalf("result = %+v, %v", result, err)
		}
		if runtime.openCount.Load() != 1 || runtime.artifactCount.Load() != 2 || runtime.toolCount.Load() != 1 || runtime.cleanupCount.Load() != 1 {
			t.Fatalf("open/artifact/tool/cleanup = %d/%d/%d/%d", runtime.openCount.Load(), runtime.artifactCount.Load(), runtime.toolCount.Load(), runtime.cleanupCount.Load())
		}
	})

	for _, terminal := range []string{"cancel", "finalizing"} {
		t.Run(terminal+" waits for artifact then cleans once", func(t *testing.T) {
			runtime := newArtifactLifecycleRuntime(true)
			controller := make(chan SingleRequestController, 1)
			executor := &channelFakeExecutor{fn: func(ctx context.Context, req SingleRequestRequest, ctrl SingleRequestController) error {
				if err := ctrl.SubmitEnvelope(testEnvelope(req.RequestID, 1, SingleRequestStatePlanning)); err != nil {
					return err
				}
				controller <- ctrl
				return ctrl.WriteInternalArtifact(ctx, SingleRequestArtifactReview, []byte("review evidence"))
			}}
			handle, err := startSingleRequestWithToolLoop(context.Background(), executor, nil, runtime, artifactLifecycleRequest(t, "request-blocked"))
			if err != nil {
				t.Fatal(err)
			}
			ctrl := <-controller
			select {
			case <-runtime.artifactStart:
			case <-time.After(2 * time.Second):
				t.Fatal("artifact operation did not start")
			}
			if terminal == "cancel" {
				handle.Cancel()
			} else {
				if err := ctrl.SubmitEnvelope(testEnvelope("request-blocked", 2, SingleRequestStateWorking)); err != nil {
					t.Fatal(err)
				}
				if err := ctrl.SubmitEnvelope(testEnvelope("request-blocked", 3, SingleRequestStateReviewing)); err != nil {
					t.Fatal(err)
				}
				final := testEnvelope("request-blocked", 4, SingleRequestStateFinalizing)
				final.Result = &SingleRequestResult{Output: "ready after artifact"}
				if err := ctrl.SubmitEnvelope(final); err != nil {
					t.Fatal(err)
				}
			}
			time.Sleep(20 * time.Millisecond)
			if runtime.cleanupCount.Load() != 0 {
				t.Fatal("cleanup ran before the in-flight artifact settled")
			}
			close(runtime.artifactGate)
			if terminal == "cancel" {
				_, waitErr := waitForExecution(t, handle)
				if !errors.Is(waitErr, ErrSingleRequestCancelled) {
					t.Fatalf("cancel error = %v", waitErr)
				}
			} else {
				waitForSingleRequestCleanup(t, handle)
				if err := handle.AcknowledgeTerminal(true); err != nil {
					t.Fatal(err)
				}
				if _, err := waitForExecution(t, handle); err != nil {
					t.Fatal(err)
				}
			}
			if runtime.openCount.Load() != 1 || runtime.artifactCount.Load() != 1 || runtime.cleanupCount.Load() != 1 {
				t.Fatalf("open/artifact/cleanup = %d/%d/%d", runtime.openCount.Load(), runtime.artifactCount.Load(), runtime.cleanupCount.Load())
			}
		})
	}

	t.Run("oversized content is rejected before open or send", func(t *testing.T) {
		runtime := newArtifactLifecycleRuntime(false)
		executor := &channelFakeExecutor{fn: func(ctx context.Context, req SingleRequestRequest, ctrl SingleRequestController) error {
			if err := ctrl.SubmitEnvelope(testEnvelope(req.RequestID, 1, SingleRequestStatePlanning)); err != nil {
				return err
			}
			return ctrl.WriteInternalArtifact(ctx, SingleRequestArtifactPlan, make([]byte, req.Binding.Workspace.Limits.MaxOutputBytes+1))
		}}
		handle, err := startSingleRequestWithToolLoop(context.Background(), executor, nil, runtime, artifactLifecycleRequest(t, "request-oversized"))
		if err != nil {
			t.Fatal(err)
		}
		_, waitErr := waitForExecution(t, handle)
		if !errors.Is(waitErr, ErrSingleRequestInternalArtifactBudget) {
			t.Fatalf("oversized error = %v", waitErr)
		}
		if runtime.openCount.Load() != 0 || runtime.artifactCount.Load() != 0 || runtime.cleanupCount.Load() != 0 {
			t.Fatalf("oversized request effects = open %d artifact %d cleanup %d", runtime.openCount.Load(), runtime.artifactCount.Load(), runtime.cleanupCount.Load())
		}
	})
}

func TestSingleRequestArtifactRequestWallClockBudgetOwnership(t *testing.T) {
	const (
		iterations  = 20
		rawSentinel = "RAW-ARTIFACT-BUDGET-SENTINEL"
	)
	for iteration := 0; iteration < iterations; iteration++ {
		runtime := newArtifactLifecycleRuntime(true)
		observer := &capturingObserver{}
		executor := &channelFakeExecutor{fn: func(ctx context.Context, req SingleRequestRequest, ctrl SingleRequestController) error {
			if err := ctrl.SubmitEnvelope(testEnvelope(req.RequestID, 1, SingleRequestStatePlanning)); err != nil {
				return err
			}
			return ctrl.WriteInternalArtifact(ctx, SingleRequestArtifactPlan, []byte(rawSentinel))
		}}
		request := artifactLifecycleRequest(t, "request-artifact-budget")
		request.Binding.Limits.WallClockMS = 30
		request.Binding.Limits.StageTimeoutMS = 30
		handle, err := startSingleRequestWithToolLoopObserved(context.Background(), executor, nil, runtime, request, observer, nil)
		if err != nil {
			close(runtime.artifactGate)
			t.Fatalf("iteration=%d StartSingleRequest: %v", iteration, err)
		}
		select {
		case <-runtime.artifactStart:
		case <-time.After(2 * time.Second):
			close(runtime.artifactGate)
			t.Fatalf("iteration=%d artifact operation did not start", iteration)
		}
		internal := handle.(*singleRequestHandle)
		select {
		case <-internal.execCtx.Done():
		case <-time.After(2 * time.Second):
			close(runtime.artifactGate)
			t.Fatalf("iteration=%d request wall-clock did not expire", iteration)
		}
		close(runtime.artifactGate)
		assertSingleRequestRequestBudgetOwnership(t, handle, observer, ErrSingleRequestInternalToolBudget, rawSentinel)
		if runtime.openCount.Load() != 1 || runtime.artifactCount.Load() != 1 || runtime.cleanupCount.Load() != 1 {
			t.Fatalf("iteration=%d open/artifact/cleanup = %d/%d/%d, want 1/1/1", iteration, runtime.openCount.Load(), runtime.artifactCount.Load(), runtime.cleanupCount.Load())
		}
	}
}