package openai

import (
	"context"
	"encoding/json"
	"errors"
	"reflect"
	"strings"
	"testing"

	edgeservice "iop/apps/edge/internal/service"
	"iop/packages/go/singlerequesttemplate"
)

type planController struct {
	binding       *edgeservice.SingleRequestBinding
	bridge        *singleRequestWorkToolBridge
	toolResult    edgeservice.InternalWorkspaceToolResult
	envelopes     []edgeservice.SingleRequestEnvelope
	kind          edgeservice.SingleRequestArtifactKind
	content       []byte
	writeAttempts int
	envelopeErr   error
	writeErr      error
}

func (c *planController) RequestID() string { return "request-1" }
func (c *planController) Binding() *edgeservice.SingleRequestBinding {
	binding := c.binding
	if binding == nil {
		binding, _ = edgeservice.NewSingleRequestBinding("virtual-model", "ws-ref", validStageBinding(), validStageBinding(), validStageBinding(), validLimits())
	}
	if binding.Workspace == nil {
		binding.Workspace = &edgeservice.SingleRequestWorkspaceBinding{Ref: "ws-ref", NodeID: "node", ConnectionGeneration: 1, OperationIDs: []string{"read", "list"}}
	}
	return binding
}
func (c *planController) Context() context.Context { return context.Background() }
func (c *planController) State() edgeservice.SingleRequestState {
	return edgeservice.SingleRequestStateAccepted
}
func (c *planController) ReadInternalArtifact(context.Context, edgeservice.SingleRequestArtifactKind) ([]byte, error) {
	return nil, errors.New("unused")
}
func (c *planController) SubmitEnvelope(e edgeservice.SingleRequestEnvelope) error {
	c.envelopes = append(c.envelopes, e)
	if c.envelopeErr != nil {
		return c.envelopeErr
	}
	if e.Stage == edgeservice.SingleRequestStateInternalTool && e.ToolCall != nil {
		result := c.toolResult.Clone()
		if result.Status == "" {
			result.Status = "success"
			result.Entries = []string{"apps", "go.mod"}
		}
		result.RequestID = e.ToolCall.RequestID
		result.StageID = e.ToolCall.StageID
		result.ToolCallID = e.ToolCall.ToolCallID
		return c.bridge.ContinueInternalTool(context.Background(), result)
	}
	return nil
}
func (c *planController) WriteInternalArtifact(_ context.Context, k edgeservice.SingleRequestArtifactKind, b []byte) error {
	c.writeAttempts++
	if c.writeErr != nil {
		return c.writeErr
	}
	c.kind = k
	c.content = append([]byte(nil), b...)
	return nil
}

func validPlanStageRequest() singleRequestPlanStageRequest {
	return singleRequestPlanStageRequest{
		RequestID:        "request-1",
		Task:             "Fix immutable task",
		StageBinding:     validStageBinding(),
		Limits:           validLimits(),
		NodeRef:          "node",
		SessionID:        "session",
		UsageAttribution: "principal",
		Sequence:         1,
	}
}

func planInspectionBody() []byte {
	return []byte(`{"id":"chat-plan-inspect","object":"chat.completion","created":1,"model":"served","choices":[{"index":0,"message":{"role":"assistant","content":null,"tool_calls":[{"id":"call-plan-list","type":"function","function":{"name":"workspace_list","arguments":"{\"relative_path\":\".\"}"}}]},"finish_reason":"tool_calls"}]}`)
}

func newPlanTestStage(t *testing.T, ctrl *planController, finalBody []byte, captured *[]edgeservice.ProviderPoolDispatchRequest) *singleRequestPlanStage {
	t.Helper()
	dispatch := matchingDispatch()
	call := 0
	provider := newSingleRequestProviderStage(&mockService{submit: func(_ context.Context, req edgeservice.ProviderPoolDispatchRequest) (*edgeservice.ProviderPoolDispatchResult, error) {
		if captured != nil {
			*captured = append(*captured, req)
		}
		call++
		body := finalBody
		if call == 1 {
			body = planInspectionBody()
		}
		return &edgeservice.ProviderPoolDispatchResult{Path: edgeservice.ProviderPoolPathTunnel, Tunnel: &mockTunnel{frames: framesFor(body)}, DispatchInfo: dispatch}, nil
	}})
	bridge := newSingleRequestWorkToolBridge()
	ctrl.bridge = bridge
	return newSingleRequestPlanStage(provider, bridge)
}

func TestSingleRequestPlanStageInspectsThenWritesArtifact(t *testing.T) {
	planJSON := `{"goal":"Inspect the target.","steps":["Step one.","Step two."],"verification":["Run focused tests."]}`
	planMD := "# Plan\n\n## Goal\nInspect the target.\n\n## Steps\n- [P1] Step one.\n- [P2] Step two.\n\n## Verification\n- Run focused tests.\n"
	var captured []edgeservice.ProviderPoolDispatchRequest
	ctrl := &planController{}
	stage := newPlanTestStage(t, ctrl, successBodyWithThoughtSignature(planJSON), &captured)
	content, err := stage.run(context.Background(), validPlanStageRequest(), ctrl)
	if err != nil {
		t.Fatal(err)
	}
	if got := string(content); got != planMD {
		t.Fatalf("content=%q, want %q", got, planMD)
	}
	if ctrl.kind != edgeservice.SingleRequestArtifactPlan || len(ctrl.envelopes) != 3 {
		t.Fatalf("controller=%+v", ctrl)
	}
	if ctrl.envelopes[0].Stage != edgeservice.SingleRequestStatePlanning || ctrl.envelopes[1].Stage != edgeservice.SingleRequestStateInternalTool || ctrl.envelopes[1].SavedStage != edgeservice.SingleRequestStatePlanning || ctrl.envelopes[2].Stage != edgeservice.SingleRequestStatePlanning {
		t.Fatalf("envelopes=%+v", ctrl.envelopes)
	}
	if len(captured) != 2 {
		t.Fatalf("provider calls=%d, want 2", len(captured))
	}
	first, _ := captured[0].Tunnel.BuildBody("gemini-3.6-flash")
	second, _ := captured[1].Tunnel.BuildBody("gemini-3.6-flash")
	if !containsAll(string(first), "Inspect the workspace before planning", "Fix immutable task", `"tool_choice":"required"`, "workspace_read", "workspace_list", `"reasoning_effort":"high"`) {
		t.Fatalf("first body=%s", first)
	}
	if strings.Contains(string(first), "workspace_write") || strings.Contains(string(first), "workspace_command") || strings.Contains(string(first), "response_format") {
		t.Fatalf("first body exposed mutation or early completion: %s", first)
	}
	if !containsAll(string(second), `"tool_choice":"auto"`, "response_format", "call-plan-list", "go.mod") {
		t.Fatalf("second body=%s", second)
	}
	var decoded map[string]any
	if err := json.Unmarshal(second, &decoded); err != nil {
		t.Fatal(err)
	}
	encodedFormat, _ := json.Marshal(singleRequestPlanResponseFormat())
	var expectedFormat any
	if err := json.Unmarshal(encodedFormat, &expectedFormat); err != nil {
		t.Fatal(err)
	}
	if got := decoded["response_format"]; !reflect.DeepEqual(got, expectedFormat) {
		t.Fatalf("response_format=%#v, want %#v", got, expectedFormat)
	}
}

func TestSingleRequestPlanStageCustomTemplate(t *testing.T) {
	customTmpl := "# Plan\n\nCustom Header\n\n## Goal\n{{goal}}\n\n## Steps\n{{steps}}\n\n## Verification\n{{verification}}\n"
	planJSON := `{"goal":"Inspect custom target.","steps":["Custom step 1.","Custom step 2."],"verification":["Custom verify."]}`
	planMD := "# Plan\n\nCustom Header\n\n## Goal\nInspect custom target.\n\n## Steps\n- [P1] Custom step 1.\n- [P2] Custom step 2.\n\n## Verification\n- Custom verify.\n"
	binding, err := edgeservice.NewSingleRequestBindingWithTemplates("virtual-model", "ws-ref", validStageBinding(), validStageBinding(), validStageBinding(), validLimits(), edgeservice.SingleRequestTemplateBinding{Plan: customTmpl, Review: singlerequesttemplate.DefaultReviewTemplate})
	if err != nil {
		t.Fatal(err)
	}
	ctrl := &planController{binding: binding}
	content, err := newPlanTestStage(t, ctrl, successBody(planJSON), nil).run(context.Background(), validPlanStageRequest(), ctrl)
	if err != nil {
		t.Fatal(err)
	}
	if string(content) != planMD {
		t.Fatalf("content=%q, want %q", string(content), planMD)
	}
}

func TestSingleRequestPlanStageRejectsFailedInspection(t *testing.T) {
	ctrl := &planController{toolResult: edgeservice.InternalWorkspaceToolResult{Status: "error", ErrorCode: "not_found"}}
	planJSON := `{"goal":"Guess.","steps":["Step one.","Step two."],"verification":["Verify."]}`
	_, err := newPlanTestStage(t, ctrl, successBody(planJSON), nil).run(context.Background(), validPlanStageRequest(), ctrl)
	if !errors.Is(err, errSingleRequestPlanStage) || ctrl.writeAttempts != 0 {
		t.Fatalf("err=%v writes=%d", err, ctrl.writeAttempts)
	}
}

func TestSingleRequestPlanStageFailsClosed(t *testing.T) {
	malformed := []string{
		"", "just plain text",
		`{"goal":"Fix bug.","steps":["Step 1","Step 2"],"verification":["Verify 1"],"extra":"no"}`,
		`{"goal":"Fix bug.","steps":["Step 1"],"verification":["Verify 1"]}`,
		`{"goal":"Fix bug.","steps":["Step 1","Step 2"],"verification":[]}`,
		`{"goal":"Line 1\nLine 2","steps":["Step 1","Step 2"],"verification":["Verify 1"]}`,
	}
	for _, raw := range malformed {
		t.Run(raw, func(t *testing.T) {
			ctrl := &planController{}
			_, err := newPlanTestStage(t, ctrl, successBody(raw), nil).run(context.Background(), validPlanStageRequest(), ctrl)
			if !errors.Is(err, errSingleRequestPlanStage) || ctrl.writeAttempts != 0 {
				t.Fatalf("err=%v writes=%d", err, ctrl.writeAttempts)
			}
		})
	}

	t.Run("initial-plan-without-inspection", func(t *testing.T) {
		dispatch := matchingDispatch()
		provider := newSingleRequestProviderStage(&mockService{submit: func(context.Context, edgeservice.ProviderPoolDispatchRequest) (*edgeservice.ProviderPoolDispatchResult, error) {
			return &edgeservice.ProviderPoolDispatchResult{Path: edgeservice.ProviderPoolPathTunnel, Tunnel: &mockTunnel{frames: framesFor(successBody(`{"goal":"Guess.","steps":["S1","S2"],"verification":["V1"]}`))}, DispatchInfo: dispatch}, nil
		}})
		bridge := newSingleRequestWorkToolBridge()
		ctrl := &planController{bridge: bridge}
		_, err := newSingleRequestPlanStage(provider, bridge).run(context.Background(), validPlanStageRequest(), ctrl)
		if !errors.Is(err, errSingleRequestPlanStage) || ctrl.writeAttempts != 0 {
			t.Fatalf("err=%v writes=%d", err, ctrl.writeAttempts)
		}
	})

	t.Run("envelope-rejection", func(t *testing.T) {
		providerCalled := false
		provider := newSingleRequestProviderStage(&mockService{submit: func(context.Context, edgeservice.ProviderPoolDispatchRequest) (*edgeservice.ProviderPoolDispatchResult, error) {
			providerCalled = true
			return nil, nil
		}})
		bridge := newSingleRequestWorkToolBridge()
		ctrl := &planController{bridge: bridge, envelopeErr: errors.New("envelope error")}
		_, err := newSingleRequestPlanStage(provider, bridge).run(context.Background(), validPlanStageRequest(), ctrl)
		if !errors.Is(err, errSingleRequestPlanStage) || providerCalled {
			t.Fatalf("err=%v providerCalled=%t", err, providerCalled)
		}
	})

	t.Run("artifact-write-failure", func(t *testing.T) {
		ctrl := &planController{writeErr: errors.New("write failure")}
		planJSON := `{"goal":"Goal text.","steps":["Step 1","Step 2"],"verification":["Verify 1"]}`
		_, err := newPlanTestStage(t, ctrl, successBody(planJSON), nil).run(context.Background(), validPlanStageRequest(), ctrl)
		if !errors.Is(err, errSingleRequestPlanStage) || ctrl.writeAttempts != 1 || len(ctrl.content) != 0 {
			t.Fatalf("err=%v attempts=%d content=%q", err, ctrl.writeAttempts, ctrl.content)
		}
	})

	t.Run("invalid-construction", func(t *testing.T) {
		ctrl := &planController{}
		if _, err := newSingleRequestPlanStage(nil, nil).run(context.Background(), validPlanStageRequest(), ctrl); !errors.Is(err, errSingleRequestPlanStage) {
			t.Fatalf("nil stage err=%v", err)
		}
	})
}

func containsAll(s string, parts ...string) bool {
	for _, p := range parts {
		if !strings.Contains(s, p) {
			return false
		}
	}
	return true
}
