package executor

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"github.com/google/uuid"
	"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
	cliproxyauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
	cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor"
	sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator"
	"github.com/tidwall/gjson"
)

const (
	claudeNativeHelperSessionID = "11111111-2222-4333-8444-555555555555"
	claudeNativeHelperUserID    = `{"device_id":"ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff","account_uuid":"bbbbbbbb-bbbb-4bbb-8bbb-bbbbbbbbbbbb","session_id":"11111111-2222-4333-8444-555555555555"}`
	claudeNativeHelperCoreBetas = "oauth-2025-04-20,interleaved-thinking-2025-05-14,redact-thinking-2026-02-12,thinking-token-count-2026-05-13,context-management-2025-06-27,prompt-caching-scope-2026-01-05"
)

func claudeNativeHelperHeaders(betas, compression string) http.Header {
	headers := http.Header{
		"Accept":            {"application/json"},
		"Accept-Encoding":   {compression},
		"Content-Type":      {"application/json"},
		"User-Agent":        {"claude-cli/2.1.258 (external, cli)"},
		"X-App":             {"cli"},
		"Anthropic-Beta":    {betas},
		"Anthropic-Version": {"2023-06-01"},
		"Anthropic-Dangerous-Direct-Browser-Access": {"true"},
		"X-Claude-Code-Session-Id":                  {claudeNativeHelperSessionID},
		"X-Client-Request-Id":                       {"66666666-7777-4888-8999-aaaaaaaaaaaa"},
		"X-Stainless-Lang":                          {"js"},
		"X-Stainless-Runtime":                       {"node"},
		"X-Stainless-Package-Version":               {"0.112.1"},
		"X-Stainless-Runtime-Version":               {"v26.3.0"},
		"X-Stainless-OS":                            {"MacOS"},
		"X-Stainless-Arch":                          {"arm64"},
		"X-Stainless-Retry-Count":                   {"0"},
		"X-Stainless-Timeout":                       {"600"},
	}
	canonical := make(http.Header, len(headers))
	for name, values := range headers {
		for _, value := range values {
			canonical.Add(name, value)
		}
	}
	return canonical
}

func claudeNativeHelperOAuthAuth(baseURL string) *cliproxyauth.Auth {
	return &cliproxyauth.Auth{
		ID: "native-helper-oauth",
		Attributes: map[string]string{
			"api_key":  "sk-ant-oat-native-helper",
			"base_url": baseURL,
		},
		Metadata: claudeOAuthTestMetadata(),
	}
}

func TestApplyClaudeHeadersPreservesCallerAsyncWithoutFingerprintOptIn(t *testing.T) {
	for _, test := range []struct {
		name      string
		confirmed bool
		profile   bool
		wantAsync string
	}{
		{name: "confirmed native", confirmed: true, wantAsync: "async"},
		{name: "unconfirmed caller default", wantAsync: "async"},
		{name: "unconfirmed caller profile", profile: true},
	} {
		t.Run(test.name, func(t *testing.T) {
			request, errRequest := http.NewRequest(http.MethodPost, "https://api.anthropic.com/v1/messages?beta=true", nil)
			if errRequest != nil {
				t.Fatal(errRequest)
			}
			incoming := http.Header{"X-Stainless-Async": {"async"}}
			auth := &cliproxyauth.Auth{Attributes: map[string]string{"api_key": "test-api-key"}}
			if test.profile {
				auth.Attributes["fingerprint_profile"] = "claude-code-cli"
			}
			if errHeaders := applyClaudeHeaders(
				request,
				auth,
				"test-api-key",
				true,
				nil,
				[]byte(`{"model":"claude-haiku-4-5-20251001"}`),
				&config.Config{},
				incoming,
				test.confirmed,
				claudeNativeHelperSessionID,
			); errHeaders != nil {
				t.Fatalf("applyClaudeHeaders() error = %v", errHeaders)
			}
			if got := request.Header.Get("X-Stainless-Async"); got != test.wantAsync {
				t.Fatalf("X-Stainless-Async = %q, want %q", got, test.wantAsync)
			}
		})
	}
}

func TestClaudeExecutorMinimalNativeHelperPreservesMarkerlessWire(t *testing.T) {
	var upstreamBody []byte
	var upstreamHeaders http.Header
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		upstreamBody, _ = io.ReadAll(r.Body)
		upstreamHeaders = r.Header.Clone()
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"id":"msg_1","type":"message","role":"assistant","model":"claude-haiku-4-5-20251001","content":[{"type":"text","text":"ok"}],"usage":{"input_tokens":1,"output_tokens":1}}`))
	}))
	defer server.Close()

	payload := []byte(`{"model":"claude-haiku-4-5-20251001","max_tokens":1,"messages":[{"role":"user","content":"helper probe"}],"metadata":{"user_id":"` + strings.ReplaceAll(claudeNativeHelperUserID, `"`, `\"`) + `"}}`)
	headers := claudeNativeHelperHeaders(claudeNativeHelperCoreBetas, "gzip, deflate, br, zstd")
	executor := NewClaudeExecutor(&config.Config{})
	_, errExecute := executor.Execute(context.Background(), claudeNativeHelperOAuthAuth(server.URL), cliproxyexecutor.Request{
		Model:   "claude-haiku-4-5-20251001",
		Payload: payload,
	}, cliproxyexecutor.Options{
		SourceFormat:    sdktranslator.FormatClaude,
		OriginalRequest: payload,
		Headers:         headers,
	})
	if errExecute != nil {
		t.Fatalf("Execute() error = %v", errExecute)
	}

	for _, path := range []string{"system", "stream", "context_management", "output_config"} {
		if got := gjson.GetBytes(upstreamBody, path); got.Exists() {
			t.Fatalf("helper body unexpectedly contains %s=%s: %s", path, got.Raw, upstreamBody)
		}
	}
	if bytes.Contains(upstreamBody, []byte(`"cache_control"`)) {
		t.Fatalf("helper body unexpectedly contains cache_control: %s", upstreamBody)
	}
	if got := gjson.GetBytes(upstreamBody, "messages.0.content").String(); got != "helper probe" {
		t.Fatalf("messages.0.content = %q, want preserved string", got)
	}
	if !bytes.HasPrefix(upstreamBody, []byte(`{"model":"claude-haiku-4-5-20251001","max_tokens":1,"messages":`)) {
		t.Fatalf("helper top-level order changed: %s", upstreamBody)
	}
	assertClaudeNativeHelperHeaders(t, upstreamHeaders, headers)
}

func TestClaudeExecutorStructuredNativeHelperPreservesStreamProfile(t *testing.T) {
	var upstreamBody []byte
	var upstreamHeaders http.Header
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		upstreamBody, _ = io.ReadAll(r.Body)
		upstreamHeaders = r.Header.Clone()
		w.Header().Set("Content-Type", "text/event-stream")
		_, _ = fmt.Fprint(w, "event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_1\",\"type\":\"message\",\"role\":\"assistant\",\"model\":\"claude-haiku-4-5-20251001\",\"content\":[],\"stop_reason\":null,\"usage\":{\"input_tokens\":1,\"output_tokens\":0}}}\n\nevent: message_stop\ndata: {\"type\":\"message_stop\"}\n\n")
	}))
	defer server.Close()

	betas := claudeNativeHelperCoreBetas + ",structured-outputs-2025-12-15"
	payload := []byte(`{"model":"claude-haiku-4-5-20251001","messages":[{"role":"user","content":[{"type":"text","text":"helper probe"}]}],"system":[{"type":"text","text":"x-anthropic-billing-header: cc_version=2.1.258; cc_entrypoint=cli; cch=00000;"},{"type":"text","text":"You are Claude Code, Anthropic's official CLI for Claude."},{"type":"text","text":"Return a short title."}],"tools":[],"metadata":{"user_id":"` + strings.ReplaceAll(claudeNativeHelperUserID, `"`, `\"`) + `"},"max_tokens":32000,"thinking":{"type":"disabled"},"temperature":1,"output_config":{"format":{"type":"json_schema","schema":{"type":"object","properties":{"title":{"type":"string"}},"required":["title"],"additionalProperties":false}}},"stream":true}`)
	headers := claudeNativeHelperHeaders(betas, "gzip, deflate, br, zstd")
	executor := NewClaudeExecutor(&config.Config{})
	result, errStream := executor.ExecuteStream(context.Background(), claudeNativeHelperOAuthAuth(server.URL), cliproxyexecutor.Request{
		Model:   "claude-haiku-4-5-20251001",
		Payload: payload,
	}, cliproxyexecutor.Options{
		SourceFormat:    sdktranslator.FormatClaude,
		OriginalRequest: payload,
		Headers:         headers,
	})
	if errStream != nil {
		t.Fatalf("ExecuteStream() error = %v", errStream)
	}
	for chunk := range result.Chunks {
		if chunk.Err != nil {
			t.Fatalf("stream chunk error = %v", chunk.Err)
		}
	}

	if got := gjson.GetBytes(upstreamBody, "system.#").Int(); got != 3 {
		t.Fatalf("system block count = %d, want native 3: %s", got, upstreamBody)
	}
	if !bytes.HasPrefix(upstreamBody, []byte(`{"model":"claude-haiku-4-5-20251001","messages":`)) {
		t.Fatalf("structured helper top-level order changed: %s", upstreamBody)
	}
	for _, path := range []string{"context_management", "output_config.effort"} {
		if got := gjson.GetBytes(upstreamBody, path); got.Exists() {
			t.Fatalf("structured helper unexpectedly contains %s=%s: %s", path, got.Raw, upstreamBody)
		}
	}
	if bytes.Contains(upstreamBody, []byte(`"cache_control"`)) {
		t.Fatalf("structured helper unexpectedly contains cache_control: %s", upstreamBody)
	}
	if got := gjson.GetBytes(upstreamBody, "stream").Bool(); !got {
		t.Fatalf("structured helper stream = false, want true: %s", upstreamBody)
	}
	if got := gjson.GetBytes(upstreamBody, "system.0.text").String(); strings.Contains(got, "cch=00000") || !strings.Contains(got, " cch=") {
		t.Fatalf("structured helper billing CCH was not re-signed: %q", got)
	}
	assertClaudeNativeHelperHeaders(t, upstreamHeaders, headers)
}

func assertClaudeNativeHelperHeaders(t *testing.T, got, incoming http.Header) {
	t.Helper()
	if got.Get("Anthropic-Beta") != incoming.Get("Anthropic-Beta") {
		t.Fatalf("Anthropic-Beta = %q, want exact native helper profile %q", got.Get("Anthropic-Beta"), incoming.Get("Anthropic-Beta"))
	}
	if strings.Contains(got.Get("Anthropic-Beta"), claudeExtendedCacheTTLBeta) || strings.Contains(got.Get("Anthropic-Beta"), claudeCodeBeta) {
		t.Fatalf("Anthropic-Beta gained standard Claude Code cache betas: %q", got.Get("Anthropic-Beta"))
	}
	for _, name := range []string{
		"Accept",
		"Accept-Encoding",
		"Content-Type",
		"User-Agent",
		"X-App",
		"Anthropic-Version",
		"Anthropic-Dangerous-Direct-Browser-Access",
		"X-Claude-Code-Session-Id",
		"X-Client-Request-Id",
		"X-Stainless-Async",
		"X-Stainless-Lang",
		"X-Stainless-Runtime",
		"X-Stainless-Package-Version",
		"X-Stainless-Runtime-Version",
		"X-Stainless-OS",
		"X-Stainless-Arch",
		"X-Stainless-Retry-Count",
		"X-Stainless-Timeout",
	} {
		gotValue := claudeNativeHelperHeaderValue(got, name)
		wantValue := claudeNativeHelperHeaderValue(incoming, name)
		if gotValue != wantValue {
			t.Fatalf("%s = %q, want preserved %q", name, gotValue, wantValue)
		}
	}
}

func claudeNativeHelperHeaderValue(headers http.Header, name string) string {
	for key, values := range headers {
		if strings.EqualFold(key, name) {
			return strings.Join(values, ",")
		}
	}
	return ""
}

// The measured minimal helper has no system field at all, so injecting a billing
// header would itself be the deviation. Keying the fallback on system presence means
// that if a payload rule later attaches a system prompt, the billing header and its
// CCH come back instead of shipping a system block native would never send unsigned.
func TestClaudeBodyNeedsBillingFallbackTracksSystemPresence(t *testing.T) {
	tests := []struct {
		name string
		body string
		want bool
	}{
		{
			name: "measured minimal helper has no system",
			body: `{"model":"claude-haiku-4-5-20251001","max_tokens":1,"messages":[{"role":"user","content":"probe"}]}`,
			want: false,
		},
		{
			name: "structured helper carries its own billing header",
			body: `{"system":[{"type":"text","text":"x-anthropic-billing-header: cc_version=2.1.258; cc_entrypoint=cli; cch=00000;"}]}`,
			want: true,
		},
		{
			name: "pipeline attached a system prompt without a billing header",
			body: `{"system":[{"type":"text","text":"injected by a payload rule"}]}`,
			want: true,
		},
		{
			name: "string system prompt also needs the fallback",
			body: `{"system":"injected by a payload rule"}`,
			want: true,
		},
	}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			if got := claudeBodyNeedsBillingFallback([]byte(test.body)); got != test.want {
				t.Fatalf("claudeBodyNeedsBillingFallback() = %v, want %v", got, test.want)
			}
		})
	}
}

// A 2.1.258 helper that reaches CPA through ANTHROPIC_BASE_URL carries no
// x-client-request-id, because the client attaches it only for a first-party base
// URL. The header must stay absent on a custom upstream instead of being
// synthesized, while a helper that did carry one keeps its own value.
func TestApplyClaudeHeadersHelperRequestIDFollowsUpstreamBase(t *testing.T) {
	const nativeRequestID = "66666666-7777-4888-8999-aaaaaaaaaaaa"
	for _, test := range []struct {
		name          string
		url           string
		incomingID    string
		wantGenerated bool
		wantID        string
	}{
		{name: "first party base without caller id", url: "https://api.anthropic.com/v1/messages?beta=true", wantGenerated: true},
		{name: "custom base without caller id", url: "https://gateway.example.com/v1/messages"},
		{name: "custom base keeps caller id", url: "https://gateway.example.com/v1/messages", incomingID: nativeRequestID, wantID: nativeRequestID},
	} {
		t.Run(test.name, func(t *testing.T) {
			request, errRequest := http.NewRequest(http.MethodPost, test.url, nil)
			if errRequest != nil {
				t.Fatal(errRequest)
			}
			incoming := http.Header{}
			if test.incomingID != "" {
				incoming.Set("X-Client-Request-Id", test.incomingID)
			}
			if errHeaders := applyClaudeHeadersWithNativeProfile(
				request,
				&cliproxyauth.Auth{Attributes: map[string]string{"api_key": "test-api-key"}},
				"test-api-key",
				false,
				nil,
				[]byte(`{"model":"claude-haiku-4-5-20251001"}`),
				&config.Config{},
				incoming,
				true,
				true,
				claudeNativeHelperSessionID,
			); errHeaders != nil {
				t.Fatalf("applyClaudeHeadersWithNativeProfile() error = %v", errHeaders)
			}
			got := request.Header.Get("X-Client-Request-Id")
			switch {
			case test.wantID != "":
				if got != test.wantID {
					t.Fatalf("X-Client-Request-Id = %q, want caller value %q", got, test.wantID)
				}
			case test.wantGenerated:
				if _, errParse := uuid.Parse(got); errParse != nil {
					t.Fatalf("X-Client-Request-Id = %q, want a generated UUID", got)
				}
			default:
				if got != "" {
					t.Fatalf("X-Client-Request-Id = %q, want the header to stay absent", got)
				}
			}
		})
	}
}