package chat

import (
	"context"
	"errors"
	"fmt"
	"net/http"
	"net/http/httptest"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/Tencent/WeKnora/internal/models/requestbudget"
	"github.com/Tencent/WeKnora/internal/types"
	secutils "github.com/Tencent/WeKnora/internal/utils"
)

type chatBudgetGuard struct {
	mu             sync.Mutex
	limit          int
	models, phases []string
	checkErr       error
}

func (g *chatBudgetGuard) Key() string { return "local-test" }
func (g *chatBudgetGuard) Check(ctx context.Context) error {
	if g.checkErr != nil {
		return g.checkErr
	}
	return ctx.Err()
}
func (g *chatBudgetGuard) Reserve(ctx context.Context, model, phase string) error {
	g.mu.Lock()
	defer g.mu.Unlock()
	if err := ctx.Err(); err != nil {
		return err
	}
	if len(g.models) >= g.limit {
		return requestbudget.ErrStopped
	}
	g.models, g.phases = append(g.models, model), append(g.phases, phase)
	return nil
}

func chatBudgetContext(g *chatBudgetGuard) context.Context {
	return requestbudget.WithModel(requestbudget.WithPhase(requestbudget.WithGuard(context.Background(), g), "wiki_generate"), "caller-forged-model")
}

func TestRequestBudgetChatSDKRawAndStreaming(t *testing.T) {
	t.Setenv("SSRF_WHITELIST", "127.0.0.1")
	secutils.ResetSSRFWhitelistForTest()
	t.Cleanup(secutils.ResetSSRFWhitelistForTest)
	for _, raw := range []bool{false, true} {
		for _, stream := range []bool{false, true} {
			for _, limit := range []int{0, 1} {
				t.Run(fmt.Sprintf("raw=%t/stream=%t/limit=%d", raw, stream, limit), func(t *testing.T) {
					var sends atomic.Int32
					server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
						sends.Add(1)
						if stream {
							w.Header().Set("Content-Type", "text/event-stream")
							fmt.Fprint(w, "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"ok\"}}]}\n\ndata: [DONE]\n\n")
						} else {
							w.Header().Set("Content-Type", "application/json")
							fmt.Fprint(w, `{"choices":[{"message":{"role":"assistant","content":"ok"},"finish_reason":"stop"}]}`)
						}
					}))
					defer server.Close()
					c, err := NewRemoteAPIChat(&ChatConfig{BaseURL: server.URL, ModelID: "concrete-chat", ModelName: "gpt-4o", APIKey: "local-only", Provider: "openai"})
					if err != nil {
						t.Fatal(err)
					}
					opts := &ChatOptions{CacheRetention: CacheRetentionNone}
					if raw {
						opts.CacheRetention, opts.PromptCacheKey = CacheRetentionShort, "force-raw"
					}
					g := &chatBudgetGuard{limit: limit}
					ctx, cancel := context.WithTimeout(chatBudgetContext(g), 3*time.Second)
					defer cancel()
					messages := []Message{{Role: "user", Content: "local test"}}
					_, _, actualRaw, err := c.buildOutbound(ctx, messages, opts, stream)
					if err != nil || actualRaw != raw {
						t.Fatalf("test route mismatch: %t %v", actualRaw, err)
					}
					if stream {
						var responses <-chan types.StreamResponse
						responses, err = c.ChatStream(ctx, messages, opts)
						if err == nil {
							for range responses {
							}
						}
					} else {
						_, err = c.Chat(ctx, messages, opts)
					}
					if limit == 0 && !errors.Is(err, requestbudget.ErrStopped) {
						t.Fatalf("expected stop: %v", err)
					}
					if limit == 1 && err != nil {
						t.Fatal(err)
					}
					if int(sends.Load()) != limit || len(g.models) != limit {
						t.Fatalf("unaccounted sends: %d models=%v", sends.Load(), g.models)
					}
					if limit == 1 && (g.models[0] != "concrete-chat" || g.phases[0] != "wiki_generate") {
						t.Fatalf("identity/phase lost: %v %v", g.models, g.phases)
					}
				})
			}
		}
	}
}

func TestRequestBudgetMultimodalFallbackCountsEachSend(t *testing.T) {
	t.Setenv("SSRF_WHITELIST", "127.0.0.1")
	secutils.ResetSSRFWhitelistForTest()
	t.Cleanup(secutils.ResetSSRFWhitelistForTest)
	for _, stream := range []bool{false, true} {
		for _, limit := range []int{1, 2} {
			t.Run(fmt.Sprintf("stream=%t/limit=%d", stream, limit), func(t *testing.T) {
				var sends atomic.Int32
				server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
					if sends.Add(1) == 1 {
						w.Header().Set("Content-Type", "application/json")
						w.WriteHeader(http.StatusBadRequest)
						fmt.Fprint(w, `{"error":{"message":"image input not supported","type":"invalid_request_error"}}`)
						return
					}
					if stream {
						w.Header().Set("Content-Type", "text/event-stream")
						fmt.Fprint(w, "data: [DONE]\n\n")
					} else {
						w.Header().Set("Content-Type", "application/json")
						fmt.Fprint(w, `{"choices":[{"message":{"role":"assistant","content":"ok"}}]}`)
					}
				}))
				defer server.Close()
				c, err := NewRemoteAPIChat(&ChatConfig{BaseURL: server.URL, ModelID: "concrete-chat", ModelName: "gpt-4o", APIKey: "local-only", Provider: "openai"})
				if err != nil {
					t.Fatal(err)
				}
				g := &chatBudgetGuard{limit: limit}
				ctx, cancel := context.WithTimeout(chatBudgetContext(g), 3*time.Second)
				defer cancel()
				messages := []Message{{Role: "user", Content: "local test", Images: []string{"data:image/png;base64,AA=="}}}
				opts := &ChatOptions{CacheRetention: CacheRetentionNone}
				if stream {
					var responses <-chan types.StreamResponse
					responses, err = c.ChatStream(ctx, messages, opts)
					if err == nil {
						for range responses {
						}
					}
				} else {
					_, err = c.Chat(ctx, messages, opts)
				}
				if limit == 1 && !errors.Is(err, requestbudget.ErrStopped) {
					t.Fatalf("fallback escaped budget: %v", err)
				}
				if limit == 2 && err != nil {
					t.Fatal(err)
				}
				if int(sends.Load()) != limit || len(g.models) != limit {
					t.Fatalf("fallback sends=%d reservations=%d", sends.Load(), len(g.models))
				}
				for i := range g.models {
					if g.models[i] != "concrete-chat" || g.phases[i] != "wiki_generate" {
						t.Fatal("fallback context lost")
					}
				}
			})
		}
	}
}

func TestRequestBudgetChatFactoryFailsClosed(t *testing.T) {
	g := &chatBudgetGuard{limit: 1}
	ctx := chatBudgetContext(g)
	c, err := NewChat(&ChatConfig{Source: types.ModelSourceLocal, ModelID: "local-model"}, nil)
	if err != nil {
		t.Fatal(err)
	}
	if _, err := c.Chat(ctx, nil, nil); !errors.Is(err, requestbudget.ErrStopped) {
		t.Fatalf("local adapter invoked: %v", err)
	}
	if _, err := c.ChatStream(ctx, nil, nil); !errors.Is(err, requestbudget.ErrStopped) {
		t.Fatalf("local stream invoked: %v", err)
	}
	if len(g.models) != 0 {
		t.Fatal("unsupported adapter reserved a send")
	}
	unknown, err := wrapChatRequestBudget(&fakeChat{id: "unknown"}, nil)
	if err != nil {
		t.Fatal(err)
	}
	if _, err := unknown.Chat(ctx, nil, nil); !errors.Is(err, requestbudget.ErrStopped) {
		t.Fatalf("unknown adapter allowed: %v", err)
	}
}

func TestRequestBudgetAnthropicFactory(t *testing.T) {
	t.Setenv("SSRF_WHITELIST", "127.0.0.1")
	secutils.ResetSSRFWhitelistForTest()
	t.Cleanup(secutils.ResetSSRFWhitelistForTest)
	for _, stream := range []bool{false, true} {
		for _, limit := range []int{0, 1} {
			t.Run(fmt.Sprintf("stream=%t/limit=%d", stream, limit), func(t *testing.T) {
				var sends atomic.Int32
				server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
					sends.Add(1)
					if stream {
						w.Header().Set("Content-Type", "text/event-stream")
						fmt.Fprint(w, "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n")
					} else {
						w.Header().Set("Content-Type", "application/json")
						fmt.Fprint(w, `{"type":"message","role":"assistant","content":[{"type":"text","text":"ok"}],"stop_reason":"end_turn"}`)
					}
				}))
				defer server.Close()
				c, err := NewChat(&ChatConfig{Source: types.ModelSourceRemote, BaseURL: server.URL, ModelID: "concrete-chat", ModelName: "claude-test", APIKey: "local-only", Provider: "anthropic"}, nil)
				if err != nil {
					t.Fatal(err)
				}
				g := &chatBudgetGuard{limit: limit}
				ctx, cancel := context.WithTimeout(chatBudgetContext(g), 3*time.Second)
				defer cancel()
				messages := []Message{{Role: "user", Content: "local test"}}
				if stream {
					var responses <-chan types.StreamResponse
					responses, err = c.ChatStream(ctx, messages, nil)
					if err == nil {
						for range responses {
						}
					}
				} else {
					_, err = c.Chat(ctx, messages, nil)
				}
				if limit == 0 && !errors.Is(err, requestbudget.ErrStopped) {
					t.Fatalf("guard not applied: %v", err)
				}
				if limit == 1 && err != nil {
					t.Fatal(err)
				}
				if int(sends.Load()) != limit || len(g.models) != limit {
					t.Fatalf("sends=%d models=%v", sends.Load(), g.models)
				}
				if limit == 1 && (g.models[0] != "concrete-chat" || g.phases[0] != "wiki_generate") {
					t.Fatalf("context lost: %v %v", g.models, g.phases)
				}
			})
		}
	}
}