package server

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

func testServer(t *testing.T) *Server {
	t.Helper()
	s, e := New(":memory:", "test-password", fstest.MapFS{"index.html": &fstest.MapFile{Data: []byte("SPA")}})
	if e != nil {
		t.Fatal(e)
	}
	t.Cleanup(func() { s.Close() })
	return s
}
func call(s *Server, m, p, b string, c *http.Cookie) *httptest.ResponseRecorder {
	r := httptest.NewRequest(m, p, strings.NewReader(b))
	r.Header.Set("Origin", "http://example.com")
	r.Header.Set("X-Requested-With", "ai-router")
	if c != nil {
		r.AddCookie(c)
	}
	w := httptest.NewRecorder()
	s.ServeHTTP(w, r)
	return w
}
func loginCookie(t *testing.T, s *Server) *http.Cookie {
	t.Helper()
	w := call(s, "POST", "/admin/api/login", `{"password":"test-password"}`, nil)
	if w.Code != 200 {
		t.Fatal(w.Code, w.Body.String())
	}
	return w.Result().Cookies()[0]
}
func TestAuthentication(t *testing.T) {
	s := testServer(t)
	for _, origin := range []string{"", "http://evil.test", "https://example.com"} {
		r := httptest.NewRequest("POST", "/admin/api/login", strings.NewReader(`{"password":"test-password"}`))
		r.Header.Set("Origin", origin)
		r.Header.Set("X-Requested-With", "ai-router")
		w := httptest.NewRecorder()
		s.ServeHTTP(w, r)
		if w.Code != 403 {
			t.Fatal(origin, w.Code)
		}
	}
	c := loginCookie(t, s)
	if !c.HttpOnly || c.SameSite != http.SameSiteStrictMode {
		t.Fatal(c)
	}
	if w := call(s, "GET", "/admin/api/channels", "", nil); w.Code != 401 {
		t.Fatal(w.Code)
	}
	if w := call(s, "POST", "/admin/api/logout", "", c); w.Code != 200 {
		t.Fatal(w.Code)
	}
	if w := call(s, "GET", "/admin/api/channels", "", c); w.Code != 401 {
		t.Fatal(w.Code)
	}
	if w := call(s, "GET", "/admin/api/session", "", nil); w.Code != 200 || w.Header().Get("Content-Security-Policy") == "" || w.Header().Get("X-Frame-Options") != "DENY" || w.Header().Get("Referrer-Policy") != "no-referrer" {
		t.Fatal("security headers missing", w.Code, w.Header())
	}
	for i := 0; i < 10; i++ {
		call(s, "POST", "/admin/api/login", `{"password":"wrong"}`, nil)
	}
	if w := call(s, "POST", "/admin/api/login", `{"password":"wrong"}`, nil); w.Code != 429 {
		t.Fatal(w.Code)
	}
	// While locked, a correct password gets the same 429: the lockout must
	// not become an oracle that confirms candidate passwords.
	if w := call(s, "POST", "/admin/api/login", `{"password":"test-password"}`, nil); w.Code != 429 {
		t.Fatal(w.Code)
	}
	// The lockout is per IP; a fresh address still logs in.
	r := httptest.NewRequest("POST", "/admin/api/login", strings.NewReader(`{"password":"test-password"}`))
	r.Header.Set("Origin", "http://example.com")
	r.Header.Set("X-Requested-With", "ai-router")
	r.RemoteAddr = "192.0.2.9:1"
	w := httptest.NewRecorder()
	s.ServeHTTP(w, r)
	if w.Code != 200 {
		t.Fatal("fresh IP rejected:", w.Code)
	}
}
func TestMappingPermissions(t *testing.T) {
	cs := []channel{{ID: "stable", Name: "renamed", Enabled: true, Models: map[string]string{"alias": "vendor/real", "vendor/real": "wrong"}}}
	k := apiKey{Channels: []string{"stable"}}
	_, actual, e := resolve(cs, k, "renamed/alias")
	if e != nil || actual != "vendor/real" {
		t.Fatal(actual, e)
	}
	if _, _, e = resolve(cs, apiKey{Channels: []string{"renamed"}}, "renamed/alias"); e == nil {
		t.Fatal("name granted access")
	}
	if _, _, e = resolve(cs, k, "renamed/vendor/unknown"); e == nil {
		t.Fatal("unknown allowed")
	}
	cs[0].Enabled = false
	if _, _, e = resolve(cs, k, "renamed/alias"); e == nil {
		t.Fatal("disabled allowed")
	}
}
func TestLegacyPassthroughIgnored(t *testing.T) {
	s := testServer(t)
	cookie := loginCookie(t, s)
	up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		t.Error("unmapped model reached upstream")
		w.WriteHeader(http.StatusBadGateway)
	}))
	defer up.Close()
	body := map[string]any{
		"name": "legacy", "base_url": up.URL, "protocol": "openai-chat",
		"enabled": true, "passthrough": true, "models": map[string]string{"alias": "vendor/real"},
	}
	save := func() channel {
		t.Helper()
		b, err := json.Marshal(body)
		if err != nil {
			t.Fatal(err)
		}
		w := call(s, "POST", "/admin/api/channels", string(b), cookie)
		if w.Code != 200 || strings.Contains(w.Body.String(), `"passthrough"`) {
			t.Fatal(w.Code, w.Body.String())
		}
		var c channel
		if err := json.Unmarshal(w.Body.Bytes(), &c); err != nil {
			t.Fatal(err)
		}
		return c
	}
	c := save()
	var stored bool
	if err := s.db.QueryRow("SELECT passthrough FROM channels WHERE id=?", c.ID).Scan(&stored); err != nil || stored {
		t.Fatal("new channel enabled passthrough", stored, err)
	}
	body["id"] = c.ID
	save()
	if err := s.db.QueryRow("SELECT passthrough FROM channels WHERE id=?", c.ID).Scan(&stored); err != nil || stored {
		t.Fatal("update enabled passthrough", stored, err)
	}
	if _, err := s.db.Exec("UPDATE channels SET passthrough=1 WHERE id=?", c.ID); err != nil {
		t.Fatal(err)
	}
	cs, err := s.loadChannels()
	if err != nil {
		t.Fatal(err)
	}
	if _, actual, err := resolve(cs, apiKey{AllChannels: true}, "legacy/alias"); err != nil || actual != "vendor/real" {
		t.Fatal(actual, err)
	}
	w := call(s, "GET", "/admin/api/channels", "", cookie)
	if w.Code != 200 || strings.Contains(w.Body.String(), `"passthrough"`) {
		t.Fatal(w.Code, w.Body.String())
	}
	w = call(s, "POST", "/admin/api/keys", `{"name":"client","all_channels":true}`, cookie)
	var key map[string]string
	if err := json.Unmarshal(w.Body.Bytes(), &key); err != nil || w.Code != 200 {
		t.Fatal(w.Code, w.Body.String(), err)
	}
	for _, path := range []string{"/v1/chat/completions", "/v1/responses", "/v1/messages"} {
		for _, model := range []string{"legacy/unknown", "legacy/vendor/real"} {
			t.Run(path+"/"+model, func(t *testing.T) {
				r := httptest.NewRequest("POST", path, strings.NewReader(`{"model":"`+model+`","messages":[]}`))
				r.Header.Set("Authorization", "Bearer "+key["key"])
				w := httptest.NewRecorder()
				s.ServeHTTP(w, r)
				if w.Code != http.StatusForbidden {
					t.Fatal(w.Code, w.Body.String())
				}
			})
		}
	}
	save()
	cs, err = s.loadChannels()
	if err != nil {
		t.Fatal(err)
	}
	if _, _, err := resolve(cs, apiKey{AllChannels: true}, "legacy/unknown"); err == nil {
		t.Fatal("updated legacy channel allowed unknown model")
	}
}

func TestKeyRevocationAndProxy(t *testing.T) {
	s := testServer(t)
	c := loginCookie(t, s)
	up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/v1/chat/completions" || r.Header.Get("Authorization") != "Bearer upstream-secret" {
			t.Error(r.URL.Path, r.Header)
		}
		var b map[string]any
		json.NewDecoder(r.Body).Decode(&b)
		if b["model"] != "vendor/real" || b["custom"] != "keep" {
			t.Error(b)
		}
		writeJSON(w, map[string]any{"choices": []any{}, "usage": map[string]int{"prompt_tokens": 7, "completion_tokens": 3}})
	}))
	defer up.Close()
	b, _ := json.Marshal(channel{Name: "demo", BaseURL: up.URL + "/v1", Protocol: "openai-chat", Enabled: true, Models: map[string]string{"alias": "vendor/real"}, APIKey: "upstream-secret"})
	w := call(s, "POST", "/admin/api/channels", string(b), c)
	if w.Code != 200 {
		t.Fatal(w.Body.String())
	}
	var ch channel
	json.Unmarshal(w.Body.Bytes(), &ch)
	if strings.Contains(w.Body.String(), "upstream-secret") {
		t.Fatal("secret leak")
	}
	w = call(s, "POST", "/admin/api/keys", `{"name":"client","channels":["`+ch.ID+`"]}`, c)
	if w.Code != 200 {
		t.Fatal(w.Body.String())
	}
	var key map[string]string
	json.Unmarshal(w.Body.Bytes(), &key)
	r := httptest.NewRequest("POST", "/v1/chat/completions", strings.NewReader(`{"model":"demo/alias","messages":[],"custom":"keep"}`))
	r.Header.Set("x-api-key", key["key"])
	w = httptest.NewRecorder()
	s.ServeHTTP(w, r)
	if w.Code != 200 {
		t.Fatal(w.Code, w.Body.String())
	}
	var id, stored string
	if e := s.db.QueryRow("SELECT id,hash FROM keys").Scan(&id, &stored); e != nil || stored == key["key"] {
		t.Fatal(e, stored)
	}
	w = call(s, "DELETE", "/admin/api/keys/"+id, "", c)
	if w.Code != 200 {
		t.Fatal(w.Body.String())
	}
	w = call(s, "PATCH", "/admin/api/keys/"+id, `{"enabled":true}`, c)
	if w.Code != 404 {
		t.Fatal("revived revoked key", w.Code)
	}
	r = httptest.NewRequest("GET", "/v1/models", nil)
	r.Header.Set("Authorization", "Bearer "+key["key"])
	w = httptest.NewRecorder()
	s.ServeHTTP(w, r)
	if w.Code != 401 {
		t.Fatal(w.Code)
	}
	var n int
	if err := s.flushLogs(); err != nil {
		t.Fatal(err)
	}
	s.db.QueryRow("SELECT count(*) FROM logs").Scan(&n)
	if n != 1 {
		t.Fatal(n)
	}
	w = call(s, "DELETE", "/admin/api/keys/"+id+"/purge", "", c)
	if w.Code != 200 {
		t.Fatal(w.Body.String())
	}
	s.db.QueryRow("SELECT count(*) FROM keys WHERE id=?", id).Scan(&n)
	if n != 0 {
		t.Fatal("purged key still stored", n)
	}
	w = call(s, "DELETE", "/admin/api/keys/"+id+"/purge", "", c)
	if w.Code != 404 {
		t.Fatal("purged key not reported missing", w.Code)
	}
	s.db.QueryRow("SELECT count(*) FROM logs").Scan(&n)
	if n != 1 {
		t.Fatal("purge removed history", n)
	}
}

func TestKeyPurgeRequiresRevocation(t *testing.T) {
	s := testServer(t)
	c := loginCookie(t, s)
	w := call(s, "POST", "/admin/api/keys", `{"name":"gone","all_channels":true}`, c)
	var key map[string]string
	if err := json.Unmarshal(w.Body.Bytes(), &key); err != nil || w.Code != 200 {
		t.Fatal(w.Code, w.Body.String(), err)
	}
	var id string
	if err := s.db.QueryRow("SELECT id FROM keys WHERE name='gone'").Scan(&id); err != nil {
		t.Fatal(err)
	}
	if w = call(s, "DELETE", "/admin/api/keys/"+id+"/purge", "", c); w.Code != 409 {
		t.Fatal("purged active key", w.Code, w.Body.String())
	}
	if w = call(s, "DELETE", "/admin/api/keys/"+id, "", c); w.Code != 200 {
		t.Fatal(w.Body.String())
	}
	if w = call(s, "DELETE", "/admin/api/keys/"+id+"/purge", "", c); w.Code != 200 {
		t.Fatal(w.Body.String())
	}
	w = call(s, "GET", "/admin/api/keys", "", c)
	if strings.Contains(w.Body.String(), `"gone"`) {
		t.Fatal("purged key still listed", w.Body.String())
	}
	r := httptest.NewRequest("GET", "/v1/models", nil)
	r.Header.Set("Authorization", "Bearer "+key["key"])
	w = httptest.NewRecorder()
	s.ServeHTTP(w, r)
	if w.Code != http.StatusUnauthorized {
		t.Fatal("purged key still authenticates", w.Code)
	}
}
func TestStatsUnknownAndFilters(t *testing.T) {
	s := testServer(t)
	c := loginCookie(t, s)
	now := time.Now().UTC().Format(time.RFC3339Nano)
	s.record(requestLog{CreatedAt: now, KeyID: "a", KeyName: "A", Status: 200, Usage: protocol.Usage{Input: 8, Output: 2, Known: true}})
	s.record(requestLog{CreatedAt: now, KeyID: "a", KeyName: "A", Status: 400})
	s.record(requestLog{CreatedAt: now, KeyID: "b", Status: 200})
	if err := s.flushLogs(); err != nil {
		t.Fatal(err)
	}
	w := call(s, "GET", "/admin/api/stats?key_id=a", "", c)
	if w.Code != 200 {
		t.Fatal(w.Body.String())
	}
	var out struct {
		Summary summary `json:"summary"`
	}
	if e := json.Unmarshal(w.Body.Bytes(), &out); e != nil {
		t.Fatal(e)
	}
	if out.Summary.Requests != 2 || out.Summary.Success != 1 || out.Summary.Unknown != 1 || out.Summary.Input == nil || *out.Summary.Input != 8 {
		t.Fatal(w.Body.String())
	}
	w = call(s, "GET", "/admin/api/logs?key_id=b", "", c)
	if !strings.Contains(w.Body.String(), `"input_tokens":null`) {
		t.Fatal(w.Body.String())
	}
	w = call(s, "GET", "/admin/api/stats?from=invalid", "", c)
	if w.Code != 400 {
		t.Fatal(w.Code)
	}
	w = call(s, "GET", "/admin/api/unknown", "", c)
	if w.Code != 404 {
		t.Fatal(w.Code)
	}
	w = call(s, "GET", "/some/spa/path", "", nil)
	if w.Code != 200 || w.Body.String() != "SPA" {
		t.Fatal(w.Code, w.Body.String())
	}
}

func TestPublicOriginCaseInsensitive(t *testing.T) {
	s := testServer(t)
	if e := s.SetPublicOrigin("https://Router.Example.com"); e != nil {
		t.Fatal(e)
	}
	r := httptest.NewRequest("POST", "/admin/api/login", strings.NewReader(`{"password":"test-password"}`))
	r.Header.Set("Origin", "https://router.example.com")
	r.Header.Set("X-Requested-With", "ai-router")
	w := httptest.NewRecorder()
	s.ServeHTTP(w, r)
	if w.Code != 200 || !w.Result().Cookies()[0].Secure {
		t.Fatal(w.Code, w.Header())
	}
}

func TestPublicOriginAndRedirect(t *testing.T) {
	s := testServer(t)
	if e := s.SetPublicOrigin("https://router.example"); e != nil {
		t.Fatal(e)
	}
	r := httptest.NewRequest("POST", "/admin/api/login", strings.NewReader(`{"password":"test-password"}`))
	r.Header.Set("Origin", "https://router.example")
	r.Header.Set("X-Requested-With", "ai-router")
	w := httptest.NewRecorder()
	s.ServeHTTP(w, r)
	if w.Code != 200 || !w.Result().Cookies()[0].Secure {
		t.Fatal(w.Code, w.Header())
	}
	reached := false
	target := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { reached = true }))
	defer target.Close()
	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { http.Redirect(w, r, target.URL, 307) }))
	defer upstream.Close()
	req, _ := http.NewRequest("POST", upstream.URL, strings.NewReader("body"))
	req.Header.Set("Authorization", "Bearer private")
	res, e := s.client.Do(req)
	if e != nil {
		t.Fatal(e)
	}
	res.Body.Close()
	if reached || res.StatusCode != 307 {
		t.Fatal("redirect followed")
	}
}