package openai

import (
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"testing"

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

func newTestRequestContext(t *testing.T, route routeDispatch, rawBody []byte) openAIRequestContext {
	t.Helper()
	return openAIRequestContext{
		r:              httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil),
		route:          route,
		rawBody:        rawBody,
		callerMetadata: map[string]string{},
		workspace:      "ws-1",
		estimate:       11,
		contextClass:   "normal",
		endpoint:       usageEndpointChatCompletions,
	}
}

// dispatchContextFixture returns a server with strict output enabled, a
// decoded request, and the caller's raw body.
func dispatchContextFixture(t *testing.T) (*Server, chatCompletionRequest, []byte) {
	t.Helper()
	srv := NewServer(config.EdgeOpenAIConf{
		Adapter:      "ollama",
		Target:       "llama-fixed",
		StrictOutput: true,
	}, &fakeRunService{}, nil)
	req := chatCompletionRequest{
		Model:    "llama3",
		Messages: []chatMessage{{Role: "user", Content: "hello"}},
	}
	return srv, req, []byte(`{"model":"llama3"}`)
}

// TestNewChatDispatchContextDirectRoute pins that the direct route resolves the
// adapter, target, and strict-output contract up front.
func TestNewChatDispatchContextDirectRoute(t *testing.T) {
	srv, req, rawBody := dispatchContextFixture(t)

	t.Run("direct route resolves adapter, target, and strict instruction", func(t *testing.T) {
		route := routeDispatch{Adapter: "ollama", Target: "llama-fixed", SessionID: "cli"}
		requestCtx := newTestRequestContext(t, route, rawBody)
		basePrompt := promptFromMessages(req.Messages)
		outputPolicy := srv.resolveOutputPolicy(basePrompt)

		dc := srv.newChatDispatchContext(requestCtx, req, basePrompt, outputPolicy)

		if dc.submitReq.Adapter != "ollama" || dc.submitReq.Target != "llama-fixed" {
			t.Fatalf("expected the direct route to resolve adapter/target, got %+v", dc.submitReq)
		}
		if dc.submitReq.ProviderPool {
			t.Fatal("expected the direct route not to be marked provider-pool")
		}
		if instruction := strictOutputContractInstruction(outputPolicy); instruction != "" && dc.prompt == basePrompt {
			t.Fatal("expected the strict-output instruction to be prepended on the direct route")
		}
		if dc.submitReq.Workspace != "ws-1" {
			t.Fatalf("expected the workspace to survive into the run request, got %q", dc.submitReq.Workspace)
		}
	})

	t.Run("run metadata carries the estimate and context class", func(t *testing.T) {
		route := routeDispatch{Adapter: "ollama", Target: "llama-fixed"}
		requestCtx := newTestRequestContext(t, route, rawBody)
		basePrompt := promptFromMessages(req.Messages)

		dc := srv.newChatDispatchContext(requestCtx, req, basePrompt, srv.resolveOutputPolicy(basePrompt))

		if dc.runMetadata["context_class"] != dc.contextClass {
			t.Fatalf("expected metadata context_class %q, got %q", dc.contextClass, dc.runMetadata["context_class"])
		}
		if dc.runMetadata["openai_model"] != "llama3" {
			t.Fatalf("expected metadata openai_model to be set, got %q", dc.runMetadata["openai_model"])
		}
		if dc.submitReq.EstimatedInputTokens != dc.estimate {
			t.Fatalf("expected the run request estimate %d to match the context, got %d", dc.estimate, dc.submitReq.EstimatedInputTokens)
		}
	})
}

// TestNewChatDispatchContextProviderPoolRoute pins that the provider-pool route
// defers adapter/target and the strict-output contract to post-admission
// dispatch, and preserves the caller's raw body for passthrough.
func TestNewChatDispatchContextProviderPoolRoute(t *testing.T) {
	srv, req, rawBody := dispatchContextFixture(t)

	t.Run("provider-pool route defers adapter, target, and strict instruction", func(t *testing.T) {
		route := routeDispatch{SessionID: "cli", ProviderPool: true}
		requestCtx := newTestRequestContext(t, route, rawBody)
		basePrompt := promptFromMessages(req.Messages)
		outputPolicy := srv.resolveOutputPolicy(basePrompt)

		dc := srv.newChatDispatchContext(requestCtx, req, basePrompt, outputPolicy)

		if dc.submitReq.Adapter != "" || dc.submitReq.Target != "" {
			t.Fatalf("expected the pool route to leave adapter/target to admission, got %+v", dc.submitReq)
		}
		if !dc.submitReq.ProviderPool {
			t.Fatal("expected the pool route to be marked provider-pool")
		}
		if dc.submitReq.ModelGroupKey != "llama3" {
			t.Fatalf("expected the model group key to reach admission, got %q", dc.submitReq.ModelGroupKey)
		}
		// Strict output is a normalized-dispatch contract; baking it into the
		// pool prompt would leak it into raw tunnel passthrough (SDD D02).
		if dc.prompt != basePrompt {
			t.Fatalf("expected the pool prompt to stay unaugmented, got %q", dc.prompt)
		}
		// The pool estimate is the ingress estimate: the prompt was not
		// re-augmented, so it must not be recomputed.
		if dc.estimate != requestCtx.estimate {
			t.Fatalf("expected the pool route to keep the ingress estimate %d, got %d", requestCtx.estimate, dc.estimate)
		}
	})

	t.Run("raw body is preserved separately from the run request", func(t *testing.T) {
		route := routeDispatch{SessionID: "cli", ProviderPool: true}
		requestCtx := newTestRequestContext(t, route, rawBody)
		basePrompt := promptFromMessages(req.Messages)

		dc := srv.newChatDispatchContext(requestCtx, req, basePrompt, srv.resolveOutputPolicy(basePrompt))

		if string(dc.rawBody) != string(rawBody) {
			t.Fatalf("expected the caller's raw body to be preserved verbatim, got %s", dc.rawBody)
		}
	})

}

// TestChatDispatchContextWithRetrySubmitDoesNotMutate pins that deriving a
// retry-bound context leaves the original context untouched.
func TestChatDispatchContextWithRetrySubmitDoesNotMutate(t *testing.T) {
	dc := newTestChatDispatchContext(httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil), chatCompletionRequest{Model: "llama3"})
	dc.retrySubmit = nil

	derived := dc.withRetrySubmit(func(ctx context.Context, req edgeservice.SubmitRunRequest) (any, error) {
		return nil, nil
	})

	if dc.retrySubmit != nil {
		t.Fatal("expected the source context to stay unmodified")
	}
	if derived.retrySubmit == nil {
		t.Fatal("expected the derived context to carry the retry function")
	}
	if derived.req.Model != dc.req.Model {
		t.Fatal("expected the derived context to keep the request")
	}
}

// TestDispatchContextsDeriveMetadataWithoutMutatingIngress pins the immutable
// boundary between ingress identity metadata and path-specific run metadata.
// Chat and Responses add different execution fields, but neither may change
// the metadata snapshot a raw provider-tunnel branch could still read.
func TestDispatchContextsDeriveMetadataWithoutMutatingIngress(t *testing.T) {
	srv, chatReq, rawBody := dispatchContextFixture(t)
	chatIngress := newTestRequestContext(t, routeDispatch{Adapter: "ollama", Target: "llama-fixed"}, rawBody)
	chatIngress.callerMetadata = map[string]string{"principal_id": "principal-1"}

	chatDC := srv.newChatDispatchContext(
		chatIngress,
		chatReq,
		promptFromMessages(chatReq.Messages),
		srv.resolveOutputPolicy(promptFromMessages(chatReq.Messages)),
	)
	if _, ok := chatIngress.callerMetadata["openai_model"]; ok {
		t.Fatal("expected chat ingress metadata to remain unchanged")
	}
	if chatDC.runMetadata["principal_id"] != "principal-1" || chatDC.runMetadata["openai_model"] != "llama3" {
		t.Fatalf("expected derived chat metadata, got %#v", chatDC.runMetadata)
	}

	responsesIngress := &responsesRequestContext{
		openAIRequestContext: openAIRequestContext{
			r:              httptest.NewRequest(http.MethodPost, "/v1/responses", nil),
			route:          routeDispatch{Adapter: "ollama", Target: "llama-fixed", SessionID: "default"},
			rawBody:        []byte(`{"model":"llama3","input":"hello"}`),
			callerMetadata: map[string]string{"principal_id": "principal-1"},
			workspace:      "ws-1",
			endpoint:       usageEndpointResponses,
		},
		envelope: responsesEnvelope{Model: "llama3"},
	}
	responsesDC, err := srv.newResponsesDispatchContext(responsesIngress, responsesRequest{
		Model: "llama3",
		Input: json.RawMessage(`"hello"`),
	})
	if err != nil {
		t.Fatalf("newResponsesDispatchContext() error = %v", err)
	}
	if _, ok := responsesIngress.callerMetadata["openai_model"]; ok {
		t.Fatal("expected responses ingress metadata to remain unchanged")
	}
	if responsesDC.submitReq.Metadata["principal_id"] != "principal-1" || responsesDC.submitReq.Metadata["openai_model"] != "llama3" {
		t.Fatalf("expected derived responses metadata, got %#v", responsesDC.submitReq.Metadata)
	}
}