package server

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"
)

type integrationKey struct{ raw, id string }

type integrationStreamRecorder struct {
	*httptest.ResponseRecorder
	flushed chan struct{}
}

func (w *integrationStreamRecorder) Flush() {
	w.ResponseRecorder.Flush()
	select {
	case w.flushed <- struct{}{}:
	default:
	}
}

func integrationAdmin(t *testing.T, s *Server, cookie *http.Cookie, method, path string, body any, out any) {
	t.Helper()
	b, err := json.Marshal(body)
	if err != nil {
		t.Fatal(err)
	}
	w := call(s, method, path, string(b), cookie)
	if w.Code != http.StatusOK {
		t.Fatalf("%s %s: %d %s", method, path, w.Code, w.Body.String())
	}
	if out != nil {
		if err := json.Unmarshal(w.Body.Bytes(), out); err != nil {
			t.Fatal(err)
		}
	}
}

func integrationChannel(t *testing.T, s *Server, cookie *http.Cookie, name, kind, base string) channel {
	t.Helper()
	if kind != "anthropic" {
		base += "/v1"
	}
	var ch channel
	integrationAdmin(t, s, cookie, "POST", "/admin/api/channels", channel{
		Name: name, Protocol: kind, BaseURL: base, Enabled: true,
		APIKey: "upstream-secret", Models: map[string]string{"gpt-5": "gpt-4", "gpt-4": "abc/gpt-4"},
	}, &ch)
	if ch.ID == "" || ch.APIKey != "" {
		t.Fatalf("invalid channel response: %+v", ch)
	}
	return ch
}

func integrationCreateKey(t *testing.T, s *Server, cookie *http.Cookie, name string, ch channel) integrationKey {
	t.Helper()
	var created struct {
		Key string `json:"key"`
	}
	integrationAdmin(t, s, cookie, "POST", "/admin/api/keys", apiKey{Name: name, Channels: []string{ch.ID}}, &created)
	var keys []apiKey
	integrationAdmin(t, s, cookie, "GET", "/admin/api/keys", nil, &keys)
	for _, k := range keys {
		if k.Name == name && created.Key != "" {
			return integrationKey{created.Key, k.ID}
		}
	}
	t.Fatal("created key missing from management API")
	return integrationKey{}
}

func integrationRequest(kind, model string, stream bool) string {
	if kind == "openai-responses" {
		return fmt.Sprintf(`{"model":%q,"input":"hello","max_output_tokens":32,"stream":%t}`, model, stream)
	}
	return fmt.Sprintf(`{"model":%q,"messages":[{"role":"user","content":"hello"}],"max_tokens":32,"stream":%t}`, model, stream)
}

func integrationPath(kind string) string {
	return map[string]string{"openai-chat": "/v1/chat/completions", "openai-responses": "/v1/responses", "anthropic": "/v1/messages"}[kind]
}

func integrationProxy(s *Server, key integrationKey, kind, model string, stream bool) *httptest.ResponseRecorder {
	r := httptest.NewRequest("POST", integrationPath(kind), strings.NewReader(integrationRequest(kind, model, stream)))
	if kind == "anthropic" {
		r.Header.Set("x-api-key", key.raw)
	} else {
		r.Header.Set("Authorization", "Bearer "+key.raw)
	}
	w := httptest.NewRecorder()
	s.ServeHTTP(w, r)
	return w
}

func integrationResponse(kind string) string {
	switch kind {
	case "openai-chat":
		return `{"id":"chat-test","object":"chat.completion","model":"gpt-4","choices":[{"index":0,"message":{"role":"assistant","content":"hello"},"finish_reason":"stop"}],"usage":{"prompt_tokens":12,"completion_tokens":3,"prompt_tokens_details":{"cached_tokens":4}}}`
	case "openai-responses":
		return `{"id":"resp-test","object":"response","model":"gpt-4","status":"completed","output":[{"type":"message","id":"msg-test","role":"assistant","status":"completed","content":[{"type":"output_text","text":"hello","annotations":[]}]}],"usage":{"input_tokens":12,"output_tokens":3,"input_tokens_details":{"cached_tokens":4}}}`
	default:
		return `{"id":"msg-test","type":"message","role":"assistant","model":"gpt-4","content":[{"type":"text","text":"hello"}],"stop_reason":"end_turn","stop_sequence":null,"usage":{"input_tokens":6,"output_tokens":3,"cache_read_input_tokens":4,"cache_creation_input_tokens":2}}`
	}
}

func integrationSSE(kind string) string {
	switch kind {
	case "openai-chat":
		return "data: " + `{"id":"chat-test","choices":[{"index":0,"delta":{"role":"assistant","content":"hello"},"finish_reason":null}]}` + "\n\ndata: " + `{"choices":[{"index":0,"delta":{},"finish_reason":"stop"}],"usage":{"prompt_tokens":12,"completion_tokens":3,"prompt_tokens_details":{"cached_tokens":4}}}` + "\n\ndata: [DONE]\n\n"
	case "openai-responses":
		return "event: response.created\ndata: " + `{"type":"response.created","response":{"id":"resp-test","status":"in_progress"}}` + "\n\n" +
			"event: response.output_item.added\ndata: " + `{"type":"response.output_item.added","output_index":0,"item":{"type":"message","role":"assistant","content":[]}}` + "\n\n" +
			"event: response.output_text.delta\ndata: " + `{"type":"response.output_text.delta","output_index":0,"content_index":0,"delta":"hello"}` + "\n\n" +
			"event: response.completed\ndata: " + `{"type":"response.completed","response":` + integrationResponse(kind) + "}\n\n"
	default:
		return "event: message_start\ndata: " + `{"type":"message_start","message":{"id":"msg-test","role":"assistant","content":[],"usage":{"input_tokens":6,"output_tokens":0,"cache_read_input_tokens":4,"cache_creation_input_tokens":2}}}` + "\n\n" +
			"event: content_block_start\ndata: " + `{"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}` + "\n\n" +
			"event: content_block_delta\ndata: " + `{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"hello"}}` + "\n\n" +
			"event: content_block_stop\ndata: " + `{"type":"content_block_stop","index":0}` + "\n\n" +
			"event: message_delta\ndata: " + `{"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":3}}` + "\n\n" +
			"event: message_stop\ndata: " + `{"type":"message_stop"}` + "\n\n"
	}
}

func integrationLogs(t *testing.T, s *Server, c *http.Cookie, key integrationKey) []requestLog {
	t.Helper()
	if err := s.flushLogs(); err != nil {
		t.Fatal(err)
	}
	var logs []requestLog
	integrationAdmin(t, s, c, "GET", "/admin/api/logs?key_id="+key.id, nil, &logs)
	for _, l := range logs {
		if l.KeyID != key.id {
			t.Fatalf("cross-key log leak: %+v", l)
		}
	}
	return logs
}

func integrationCount(t *testing.T, name string, got *int64, want int64) {
	t.Helper()
	if got == nil || *got != want {
		t.Fatalf("%s: got %v, want %d", name, got, want)
	}
}

func TestIntegrationProtocolMatrix(t *testing.T) {
	kinds := []string{"openai-chat", "openai-responses", "anthropic"}
	for _, upKind := range kinds {
		for _, downKind := range kinds {
			t.Run(downKind+"_to_"+upKind, func(t *testing.T) {
				s := testServer(t)
				cookie := loginCookie(t, s)
				var calls atomic.Int64
				var expectedModel atomic.Value
				up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
					calls.Add(1)
					if r.URL.Path != integrationPath(upKind) {
						t.Errorf("upstream path: %s", r.URL.Path)
					}
					if upKind == "anthropic" {
						if r.Header.Get("x-api-key") != "upstream-secret" || r.Header.Get("anthropic-version") == "" {
							t.Error("missing Anthropic credentials/version")
						}
					} else if r.Header.Get("Authorization") != "Bearer upstream-secret" {
						t.Error("wrong upstream credentials")
					}
					var body struct {
						Model  string `json:"model"`
						Stream bool   `json:"stream"`
					}
					b, err := io.ReadAll(r.Body)
					if err != nil || json.Unmarshal(b, &body) != nil {
						t.Error("invalid upstream request")
						w.WriteHeader(400)
						return
					}
					if body.Model != expectedModel.Load() || !strings.Contains(string(b), "hello") {
						t.Errorf("upstream request: %s; want model %s", b, expectedModel.Load())
					}
					if body.Stream {
						w.Header().Set("Content-Type", "text/event-stream")
						for _, frame := range strings.Split(integrationSSE(upKind), "\n\n") {
							if frame != "" {
								fmt.Fprint(w, frame+"\n\n")
								w.(http.Flusher).Flush()
							}
						}
					} else {
						w.Header().Set("Content-Type", "application/json")
						fmt.Fprint(w, integrationResponse(upKind))
					}
				}))
				defer up.Close()
				abc := integrationChannel(t, s, cookie, "abc", upKind, up.URL)
				other := integrationChannel(t, s, cookie, "other", upKind, up.URL)
				keys := []integrationKey{integrationCreateKey(t, s, cookie, "first", abc), integrationCreateKey(t, s, cookie, "second", other)}
				for i, key := range keys {
					name, forbidden := "abc", "other"
					if i == 1 {
						name, forbidden = forbidden, name
					}
					for _, stream := range []bool{false, true} {
						for _, alias := range []string{"gpt-5", "gpt-4"} {
							expectedModel.Store("gpt-4")
							if alias == "gpt-4" {
								expectedModel.Store("abc/gpt-4")
							}
							w := integrationProxy(s, key, downKind, name+"/"+alias, stream)
							if w.Code != 200 || !strings.Contains(w.Body.String(), "hello") {
								t.Fatalf("stream=%t model=%s: %d %s", stream, alias, w.Code, w.Body.String())
							}
							if stream {
								terminal := map[string]string{"openai-chat": "data: [DONE]", "openai-responses": "event: response.completed", "anthropic": "event: message_stop"}[downKind]
								if !w.Flushed || !strings.Contains(w.Header().Get("Content-Type"), "text/event-stream") || !strings.Contains(w.Body.String(), terminal) {
									t.Fatalf("incomplete SSE: %s", w.Body.String())
								}
							} else {
								var body map[string]any
								if err := json.Unmarshal(w.Body.Bytes(), &body); err != nil {
									t.Fatal(err)
								}
								field := map[string]string{"openai-chat": "choices", "openai-responses": "output", "anthropic": "content"}[downKind]
								if body[field] == nil {
									t.Fatalf("wrong downstream shape: %s", w.Body.String())
								}
							}
						}
					}
					before := calls.Load()
					if w := integrationProxy(s, key, downKind, forbidden+"/gpt-5", false); w.Code != 403 {
						t.Fatalf("permission: %d %s", w.Code, w.Body.String())
					}
					if calls.Load() != before {
						t.Fatal("forbidden request reached upstream")
					}
				}
				if calls.Load() != 8 {
					t.Fatalf("upstream calls: %d", calls.Load())
				}
				for _, key := range keys {
					logs := integrationLogs(t, s, cookie, key)
					if len(logs) != 5 {
						t.Fatalf("logs: %+v", logs)
					}
					for _, l := range logs {
						if l.Status == 403 {
							if l.UsageKnown || l.Input != nil {
								t.Fatal("denied request has usage")
							}
							continue
						}
						actual := "gpt-4"
						if strings.HasSuffix(l.RequestModel, "/gpt-4") {
							actual = "abc/gpt-4"
						}
						if l.Status != 200 || l.Error != "" || !l.UsageKnown || l.ActualModel != actual || l.Downstream != downKind || l.Upstream != upKind {
							t.Fatalf("wrong log: %+v", l)
						}
						integrationCount(t, "input", l.Input, 12)
						integrationCount(t, "output", l.Output, 3)
						integrationCount(t, "cache read", l.CacheRead, 4)
						write := int64(0)
						if upKind == "anthropic" {
							write = 2
						}
						integrationCount(t, "cache write", l.CacheWrite, write)
					}
					var stats struct {
						Summary summary `json:"summary"`
						ByKey   []struct {
							KeyID string `json:"key_id"`
						} `json:"by_key"`
					}
					integrationAdmin(t, s, cookie, "GET", "/admin/api/stats?key_id="+key.id, nil, &stats)
					if stats.Summary.Requests != 5 || stats.Summary.Success != 4 || stats.Summary.Failed != 1 || stats.Summary.Unknown != 1 || len(stats.ByKey) != 1 || stats.ByKey[0].KeyID != key.id {
						t.Fatalf("wrong isolated stats: %+v", stats)
					}
					integrationCount(t, "total input", stats.Summary.Input, 48)
					integrationCount(t, "total output", stats.Summary.Output, 12)
					integrationCount(t, "total cache read", stats.Summary.CacheRead, 16)
					write := int64(0)
					if upKind == "anthropic" {
						write = 8
					}
					integrationCount(t, "total cache write", stats.Summary.CacheWrite, write)
				}
			})
		}
	}
}

func TestIntegrationStreamFailures(t *testing.T) {
	for _, tc := range []struct {
		name, body string
		status     int
	}{
		{"upstream_http_error", `{"error":"private upstream detail"}`, 429},
		{"error_event", "event: error\ndata: {\"type\":\"error\",\"error\":{\"message\":\"private upstream detail\"}}\n\n", 200},
		{"malformed_event", "data: {invalid}\n\n", 200},
		{"truncated_stream", "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"partial\"},\"finish_reason\":null}]}\n\n", 200},
	} {
		t.Run(tc.name, func(t *testing.T) {
			s := testServer(t)
			cookie := loginCookie(t, s)
			up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				w.Header().Set("Content-Type", "text/event-stream")
				w.WriteHeader(tc.status)
				fmt.Fprint(w, tc.body)
			}))
			defer up.Close()
			ch := integrationChannel(t, s, cookie, "abc", "openai-chat", up.URL)
			key := integrationCreateKey(t, s, cookie, "errors", ch)
			w := integrationProxy(s, key, "openai-responses", "abc/gpt-5", true)
			if w.Code != tc.status {
				t.Fatalf("status: %d body: %s", w.Code, w.Body.String())
			}
			if strings.Contains(w.Body.String(), "response.completed") || strings.Contains(w.Body.String(), "private upstream detail") {
				t.Fatalf("false success or leaked upstream detail: %s", w.Body.String())
			}
			logs := integrationLogs(t, s, cookie, key)
			if len(logs) != 1 {
				t.Fatalf("logs: %+v", logs)
			}
			l := logs[0]
			if l.Status < 400 || l.Error == "" || l.UsageKnown || l.Input != nil || l.Output != nil {
				t.Fatalf("failure not recorded correctly: %+v", l)
			}
		})
	}
}

func TestAnthropicHeaderForwarding(t *testing.T) {
	s := testServer(t)
	cookie := loginCookie(t, s)
	var mu sync.Mutex
	got := map[string]string{}
	up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		mu.Lock()
		got["version"], got["beta"] = r.Header.Get("anthropic-version"), r.Header.Get("anthropic-beta")
		mu.Unlock()
		w.Header().Set("Content-Type", "application/json")
		fmt.Fprint(w, integrationResponse("anthropic"))
	}))
	defer up.Close()
	ch := integrationChannel(t, s, cookie, "abc", "anthropic", up.URL)
	key := integrationCreateKey(t, s, cookie, "beta", ch)
	request := func(version, beta string) {
		r := httptest.NewRequest("POST", "/v1/messages", strings.NewReader(integrationRequest("anthropic", "abc/gpt-5", false)))
		r.Header.Set("x-api-key", key.raw)
		if version != "" {
			r.Header.Set("anthropic-version", version)
		}
		if beta != "" {
			r.Header.Set("anthropic-beta", beta)
		}
		w := httptest.NewRecorder()
		s.ServeHTTP(w, r)
		if w.Code != 200 {
			t.Fatalf("status %d: %s", w.Code, w.Body.String())
		}
	}
	request("2023-06-01", "context-1m-2025-08-07,interleaved-thinking-2025-05-14")
	mu.Lock()
	if got["version"] != "2023-06-01" || got["beta"] != "context-1m-2025-08-07,interleaved-thinking-2025-05-14" {
		t.Fatalf("valid headers dropped: %v", got)
	}
	mu.Unlock()
	request("bad version; injected", `weird"value`)
	mu.Lock()
	if got["version"] != "2023-06-01" || got["beta"] != "" {
		t.Fatalf("unsafe headers forwarded: %v", got)
	}
	mu.Unlock()
}

func TestIntegrationCancellationClosesUpstream(t *testing.T) {
	s := testServer(t)
	cookie := loginCookie(t, s)
	started, stopped := make(chan struct{}), make(chan struct{})
	release := make(chan struct{})
	up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		io.Copy(io.Discard, r.Body)
		w.Header().Set("Content-Type", "text/event-stream")
		fmt.Fprint(w, "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"partial\"},\"finish_reason\":null}]}\n\n")
		w.(http.Flusher).Flush()
		close(started)
		select {
		case <-r.Context().Done():
			close(stopped)
		case <-release:
		}
	}))
	defer up.Close()
	defer close(release)
	ch := integrationChannel(t, s, cookie, "abc", "openai-chat", up.URL)
	key := integrationCreateKey(t, s, cookie, "cancel", ch)
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	r := httptest.NewRequest("POST", "/v1/responses", strings.NewReader(integrationRequest("openai-responses", "abc/gpt-5", true))).WithContext(ctx)
	r.Header.Set("Authorization", "Bearer "+key.raw)
	done := make(chan struct{})
	w := &integrationStreamRecorder{ResponseRecorder: httptest.NewRecorder(), flushed: make(chan struct{}, 1)}
	go func() { s.ServeHTTP(w, r); close(done) }()
	select {
	case <-started:
	case <-time.After(3 * time.Second):
		t.Fatal("upstream did not start")
	}
	select {
	case <-w.flushed:
	case <-time.After(3 * time.Second):
		t.Fatal("downstream did not receive streaming output")
	}
	cancel()
	select {
	case <-done:
	case <-time.After(3 * time.Second):
		t.Fatal("proxy did not stop on cancellation")
	}
	select {
	case <-stopped:
	case <-time.After(3 * time.Second):
		t.Fatal("upstream request was not cancelled")
	}
	logs := integrationLogs(t, s, cookie, key)
	if len(logs) != 1 || logs[0].Status < 400 || logs[0].Error == "" || logs[0].UsageKnown || logs[0].Input != nil {
		t.Fatalf("cancel log: %+v", logs)
	}
	if strings.Contains(w.Body.String(), "response.completed") {
		t.Fatal("cancelled response marked completed")
	}
}