package server

import (
	"ai-router/internal/protocol"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"
)

const probePath = "/admin/api/models/test"
const probeChat = `{"id":"c","object":"chat.completion","choices":[{"message":{"role":"assistant","content":"OK"},"finish_reason":"stop"}]}`

func saveProbeChannel(t *testing.T, s *Server, cookie *http.Cookie, url, kind string, enabled bool) {
	t.Helper()
	b, _ := json.Marshal(channel{Name: "saved", BaseURL: url, Protocol: kind, Enabled: enabled, Models: map[string]string{"alias": "actual/model"}, APIKey: "secret-key"})
	if w := call(s, "POST", "/admin/api/channels", string(b), cookie); w.Code != 200 {
		t.Fatal(w.Body.String())
	}
}

func probeResult(t *testing.T, w *httptest.ResponseRecorder) modelTestResult {
	t.Helper()
	var result modelTestResult
	if w.Code != 200 || json.Unmarshal(w.Body.Bytes(), &result) != nil {
		t.Fatal(w.Code, w.Body.String())
	}
	if result.DurationMS < 0 || strings.Contains(w.Body.String(), "secret-key") {
		t.Fatal(w.Body.String())
	}
	return result
}

func TestModelProbeProtocols(t *testing.T) {
	for _, tc := range []struct{ kind, path, limit, response string }{
		{"openai-chat", "/chat/completions", "max_completion_tokens", probeChat},
		{"openai-responses", "/responses", "max_output_tokens", `{"id":"r","object":"response","status":"completed","output":[{"type":"message","role":"assistant","content":[{"type":"output_text","text":"OK"}]}]}`},
		{"anthropic", "/v1/messages", "max_tokens", `{"id":"m","type":"message","role":"assistant","stop_reason":"end_turn","content":[{"type":"text","text":"OK"}]}`},
	} {
		t.Run(tc.kind, func(t *testing.T) {
			s := testServer(t)
			cookie := loginCookie(t, s)
			expectedPrompt := "Reply with OK."
			up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				b, _ := io.ReadAll(r.Body)
				var obj map[string]any
				if json.Unmarshal(b, &obj) != nil || obj["model"] != "actual/model" || obj[tc.limit] != float64(128) || obj["stream"] != false || !strings.Contains(string(b), expectedPrompt) {
					t.Error(string(b))
				}
				if tc.kind == "openai-chat" && obj["max_tokens"] != nil {
					t.Error("legacy token limit")
				}
				if r.Method != "POST" || r.URL.Path != tc.path {
					t.Error(r.Method, r.URL.Path)
				}
				if tc.kind == "anthropic" {
					if r.Header.Get("x-api-key") != "secret-key" || r.Header.Get("anthropic-version") != "2023-06-01" {
						t.Error("headers")
					}
				} else if r.Header.Get("Authorization") != "Bearer secret-key" {
					t.Error("authorization")
				}
				fmt.Fprint(w, tc.response)
			}))
			defer up.Close()
			saveProbeChannel(t, s, cookie, up.URL, tc.kind, true)
			for _, prompt := range []string{"", strings.Repeat("中", 2000)} {
				input := map[string]string{"model": "saved/alias"}
				if prompt != "" {
					expectedPrompt = prompt
					input["prompt"] = prompt
				}
				b, _ := json.Marshal(input)
				result := probeResult(t, call(s, "POST", probePath, string(b), cookie))
				if !result.OK || result.Status != 200 || result.Error != "" {
					t.Fatal(result)
				}
			}
			for _, table := range []string{"keys", "logs"} {
				var n int
				if e := s.db.QueryRow("SELECT count(*) FROM " + table).Scan(&n); e != nil || n != 0 {
					t.Fatal(table, n, e)
				}
			}
			if len(s.logQueue) != 0 {
				t.Fatal("queued formal log")
			}
		})
	}
}

func TestModelProbeGuardAndInput(t *testing.T) {
	s := testServer(t)
	cookie := loginCookie(t, s)
	for _, prompt := range []string{"", " ", "\t\r\n", "\u3000\u00a0"} {
		body, _ := json.Marshal(map[string]string{"model": "saved/alias", "prompt": prompt})
		if w := call(s, "POST", probePath, string(body), cookie); w.Code != 400 {
			t.Fatalf("blank prompt %q: status %d", prompt, w.Code)
		}
	}
	up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { t.Error("unexpected upstream request") }))
	defer up.Close()
	saveProbeChannel(t, s, cookie, up.URL, "openai-chat", false)
	if w := call(s, "POST", probePath, `{"model":"saved/alias"}`, nil); w.Code != 401 {
		t.Fatal(w.Code)
	}
	for _, header := range []string{"Origin", "X-Requested-With"} {
		r := httptest.NewRequest("POST", probePath, strings.NewReader(`{"model":"saved/alias"}`))
		r.AddCookie(cookie)
		r.Header.Set("Origin", "http://example.com")
		r.Header.Set("X-Requested-With", "ai-router")
		r.Header.Del(header)
		w := httptest.NewRecorder()
		s.ServeHTTP(w, r)
		if w.Code != 403 {
			t.Fatal(w.Code)
		}
	}
	for _, body := range []string{`{"model":"saved/alias"}`, `{"model":"missing/alias"}`} {
		if w := call(s, "POST", probePath, body, cookie); w.Code != 403 {
			t.Fatal(w.Code)
		}
	}
	if _, e := s.db.Exec("UPDATE channels SET enabled=1"); e != nil {
		t.Fatal(e)
	}
	if w := call(s, "POST", probePath, `{"model":"saved/unmapped"}`, cookie); w.Code != 403 {
		t.Fatal(w.Code)
	}
	for _, body := range []string{`{}`, `null`, `{"model":"alias"}`, `{"model":"/alias"}`, `{"model":"saved/"}`, `{"model":"saved/alias","prompt":2}`, `{"model":"saved/alias"} {}`, `{"model":"saved/alias","prompt":"` + strings.Repeat("中", 2001) + `"}`} {
		if w := call(s, "POST", probePath, body, cookie); w.Code != 400 {
			t.Fatal(body, w.Code)
		}
	}
	if w := call(s, "GET", probePath, "", cookie); w.Code != 405 {
		t.Fatal(w.Code)
	}
}

func TestModelProbeFailures(t *testing.T) {
	for _, tc := range []struct {
		name, body string
		status     int
		error      string
	}{
		{"rejected", "secret-key upstream body", 401, "upstream rejected request"},
		{"redirect", "secret-key", 302, "upstream rejected request"},
		{"html", "<html>secret-key</html>", 200, "invalid upstream response"},
		{"empty", "{}", 200, "invalid upstream response"},
		{"oversize", strings.Repeat("x", modelTestMaxBody+1), 200, "upstream response too large"},
		{"limit", probeChat + strings.Repeat(" ", modelTestMaxBody-len(probeChat)), 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) {
				if r.URL.Path != "/chat/completions" {
					t.Error("followed redirect")
				}
				w.Header().Set("Location", "/redirected")
				w.WriteHeader(tc.status)
				fmt.Fprint(w, tc.body)
			}))
			defer up.Close()
			saveProbeChannel(t, s, cookie, up.URL, "openai-chat", true)
			result := probeResult(t, call(s, "POST", probePath, `{"model":"saved/alias"}`, cookie))
			if result.Status != tc.status || result.Error != tc.error || result.OK != (tc.error == "") {
				t.Fatal(result)
			}
		})
	}
}

func TestModelProbeResponseValidation(t *testing.T) {
	if !validModelTestResponse(protocol.Anthropic, []byte(`{"id":"m","type":"message","role":"assistant","stop_reason":"max_tokens","content":[{"type":"thinking","thinking":"reasoning","signature":"sig"}]}`)) {
		t.Fatal("valid thinking-only response rejected")
	}
	for _, tc := range []struct {
		kind protocol.Kind
		body string
		ok   bool
	}{
		{protocol.Chat, `{"id":"c","object":"chat.completion","choices":[{"message":{"role":"assistant","content":null},"finish_reason":"length"}]}`, true},
		{protocol.Responses, `{"id":"r","object":"response","status":"incomplete","incomplete_details":{"reason":"max_output_tokens"},"output":[{"type":"reasoning","summary":[]}]}`, true},
		{protocol.Anthropic, `{"id":"m","type":"message","role":"assistant","stop_reason":"max_tokens","content":[]}`, true},
		{protocol.Chat, `{"id":"c","object":"chat.completion","choices":[{}]}`, false},
		{protocol.Responses, `{"id":"r","object":"response","status":"failed","output":[]}`, false},
		{protocol.Responses, `{"id":"r","object":"response","status":"completed","output":[{}]}`, false},
		{protocol.Anthropic, `{"id":"m","type":"message","role":"assistant","stop_reason":"end_turn","content":[{}]}`, false},
	} {
		if got := validModelTestResponse(tc.kind, []byte(tc.body)); got != tc.ok {
			t.Error(tc.kind, tc.body, got)
		}
	}
	for _, kind := range []protocol.Kind{protocol.Chat, protocol.Responses, protocol.Anthropic} {
		for _, body := range []string{"{}", "null", "[]", probeChat + "{}", `{"error":{"message":"secret-key"}}`} {
			if validModelTestResponse(kind, []byte(body)) {
				t.Error(kind, body)
			}
		}
	}
}

type probeRoundTripper func(*http.Request) (*http.Response, error)

func (f probeRoundTripper) RoundTrip(r *http.Request) (*http.Response, error) { return f(r) }

func TestModelProbeTimeoutAndCancellation(t *testing.T) {
	for _, cancelRequest := range []bool{false, true} {
		t.Run(fmt.Sprint(cancelRequest), func(t *testing.T) {
			s := testServer(t)
			cookie := loginCookie(t, s)
			up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}))
			defer up.Close()
			saveProbeChannel(t, s, cookie, up.URL, "openai-chat", true)
			ctx, cancel := context.WithCancel(context.Background())
			defer cancel()
			s.client.Transport = probeRoundTripper(func(r *http.Request) (*http.Response, error) {
				deadline, ok := r.Context().Deadline()
				if !ok || time.Until(deadline) > 30*time.Second || time.Until(deadline) < 29*time.Second {
					t.Error("missing 30s deadline", deadline)
				}
				if cancelRequest {
					cancel()
					<-r.Context().Done()
					return nil, r.Context().Err()
				}
				return nil, context.DeadlineExceeded
			})
			r := httptest.NewRequest("POST", probePath, strings.NewReader(`{"model":"saved/alias"}`)).WithContext(ctx)
			r.AddCookie(cookie)
			r.Header.Set("Origin", "http://example.com")
			r.Header.Set("X-Requested-With", "ai-router")
			w := httptest.NewRecorder()
			s.ServeHTTP(w, r)
			result := probeResult(t, w)
			if result.OK || result.Status != 0 || result.Error == "" {
				t.Fatal(result)
			}
		})
	}
	s := testServer(t)
	canceled := make(chan struct{})
	up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(200)
		w.(http.Flusher).Flush()
		<-r.Context().Done()
		close(canceled)
	}))
	defer up.Close()
	ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
	defer cancel()
	result := s.probeModel(ctx, channel{BaseURL: up.URL, Protocol: "openai-chat"}, "actual", "test")
	if result.OK || result.Status != 200 || result.Error != "model test timed out" {
		t.Fatal(result)
	}
	select {
	case <-canceled:
	case <-time.After(time.Second):
		t.Fatal("upstream was not canceled")
	}
}