package server

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"testing"

	"git.misaka.ren/M1saka/zhanlu_proxy/internal/auth"
	"git.misaka.ren/M1saka/zhanlu_proxy/internal/config"
	"git.misaka.ren/M1saka/zhanlu_proxy/internal/store"
)

const testSM2Key = "8d6ee90b3c4d299ae5abd655dbc3547c110ae8aeff1de18b0df241f215f90748"

// setupTestServer spins up a mock Zhanlu upstream and a proxy server wired to
// it. The proxy is backed by a temp SQLite store so credentials persist in the
// db during the test instead of a JSON file.
func setupTestServer(t *testing.T) (*httptest.Server, *httptest.Server, *store.Store) {
	t.Helper()
	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.URL.Path {
		case "/api/query/acepilot-h5/manager/code/getAuthCode":
			writeJSON(w, http.StatusOK, map[string]any{"state": "OK"})
		case "/api/query/acepilot-h5/manager/code/checkCode":
			writeJSON(w, http.StatusOK, map[string]any{"state": "OK", "body": map[string]any{
				"result": true,
				"ak":     "BASE64AK", "sk": "BASE64SK", "license": "BASE64TOKEN",
			}})
		case "/api/acepilot/zhanlu/v1/login":
			if r.Header.Get("plugin_type") != "zhanlu_ide" {
				writeJSON(w, http.StatusBadRequest, map[string]any{"state": "ERROR", "errorMessage": "bad plugin_type"})
				return
			}
			writeJSON(w, http.StatusOK, map[string]any{"state": "OK", "body": map[string]any{
				"email": "dev@example.com", "organization": "cmcc", "team": "ai",
			}})
		case "/user/api/v2/external/key/get-or-create":
			writeJSON(w, http.StatusOK, map[string]any{"apiKey": "sk-test-456"})
		case "/chat/completions":
			if r.Header.Get("Authorization") != "Bearer sk-test-456" {
				writeJSON(w, http.StatusUnauthorized, map[string]any{"error": map[string]any{"message": "bad auth"}})
				return
			}
			w.Header().Set("Content-Type", "text/event-stream")
			_, _ = fmt.Fprint(w, "data: {\"id\":\"chatcmpl-x\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"hello\"},\"finish_reason\":null}],\"usage\":{\"prompt_tokens\":10,\"completion_tokens\":20,\"total_tokens\":30,\"prompt_tokens_details\":{\"cached_tokens\":4}}}\n\n")
			_, _ = fmt.Fprint(w, "data: {\"id\":\"chatcmpl-x\",\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n")
			_, _ = fmt.Fprint(w, "data: [DONE]\n\n")
		case "/gateway/v1/model/info":
			writeJSON(w, http.StatusOK, map[string]any{"data": []map[string]any{
				{"model_name": "GLM-4.7"}, {"id": "MiniMaxAI/MiniMax-M2.5"},
			}})
		default:
			http.NotFound(w, r)
		}
	}))

	st, err := store.Open(filepath.Join(t.TempDir(), "stats.db"))
	if err != nil {
		t.Fatalf("open store: %v", err)
	}
	t.Cleanup(func() { st.Close() })
	cfg := config.Config{
		ListenAddr:         ":0",
		MobileLoginBaseURL: upstream.URL,
		MobileModelBaseURL: upstream.URL,
		UpstreamPath:       "/chat/completions",
		DBPath:             filepath.Join(t.TempDir(), "zhanlu.db"),
		TokenDecryptKey:    "3jw7woww2rvhla6k",
		PublicKeyPEM:       defaultTestPublicKey,
		PhonePublicKeyPEM:  defaultTestPublicKey,
		SM2PrivateKey:      testSM2Key,
		PluginVersion:      "1.4.2",
	}
	h := New(cfg, st)
	proxy := httptest.NewServer(h)
	return upstream, proxy, st
}

// TestPhoneLoginAndChat exercises the full v1.4.2 flow: SMS login, profile
// fetch, SM2 API-key provisioning, then OpenAI-compatible chat and models.
func TestPhoneLoginAndChat(t *testing.T) {
	upstream, proxy, st := setupTestServer(t)
	defer upstream.Close()
	defer proxy.Close()

	// 1. request phone code
	resp, err := http.Post(proxy.URL+"/api/auth/code", "application/json", strings.NewReader(`{"telephone":"13800000000"}`))
	if err != nil {
		t.Fatal(err)
	}
	var codeResp map[string]any
	_ = json.NewDecoder(resp.Body).Decode(&codeResp)
	resp.Body.Close()
	secret, _ := codeResp["secret"].(string)

	// 2. login with phone code (server RSA-encrypts the telephone itself)
	loginBody, _ := json.Marshal(map[string]string{"telephone": "13800000000", "code": "123456", "secret": secret})
	resp, err = http.Post(proxy.URL+"/api/auth/login", "application/json", bytes.NewReader(loginBody))
	if err != nil {
		t.Fatal(err)
	}
	var loginResp map[string]any
	_ = json.NewDecoder(resp.Body).Decode(&loginResp)
	resp.Body.Close()
	if !okValue(loginResp) {
		t.Fatalf("login failed: %v", loginResp)
	}

	// 3. credentials store should contain the provisioned api key
	creds, err := st.LoadCredentials()
	if err != nil {
		t.Fatal(err)
	}
	if creds.APIKey != "sk-test-456" {
		t.Fatalf("apiKey = %q", creds.APIKey)
	}
	if creds.Email != "dev@example.com" {
		t.Fatalf("email = %q", creds.Email)
	}

	// 4. non-streaming chat completion
	chatBody := `{"model":"GLM-4.7","messages":[{"role":"user","content":"hi"}],"stream":false}`
	resp, err = http.Post(proxy.URL+"/v1/chat/completions", "application/json", strings.NewReader(chatBody))
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()
	var chatResp map[string]any
	if err := json.NewDecoder(resp.Body).Decode(&chatResp); err != nil {
		t.Fatal(err)
	}
	choices, _ := chatResp["choices"].([]any)
	if len(choices) != 1 {
		t.Fatalf("chat choices = %v", chatResp)
	}
	msg, _ := choices[0].(map[string]any)["message"].(map[string]any)
	if msg["content"] != "hello" {
		t.Fatalf("chat content = %v", msg)
	}

	// 5. models endpoint should prefer gateway model info
	resp, err = http.Get(proxy.URL + "/v1/models")
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()
	var modelsResp map[string]any
	_ = json.NewDecoder(resp.Body).Decode(&modelsResp)
	items, _ := modelsResp["data"].([]any)
	if len(items) != 2 {
		t.Fatalf("models = %v", modelsResp)
	}
}

func TestStreamingChat(t *testing.T) {
	upstream, proxy, st := setupTestServer(t)
	defer upstream.Close()
	defer proxy.Close()

	// Seed credentials directly with the api key into the store
	creds := auth.Credentials{
		AccessKey:    "AK",
		SecretKey:    "SK",
		Token:        "TOKEN",
		APIKey:       "sk-test-456",
		ModelBaseURL: upstream.URL,
		Email:        "dev@example.com",
	}
	if err := st.SaveCredentials(creds); err != nil {
		t.Fatal(err)
	}

	chatBody := `{"model":"GLM-4.7","messages":[{"role":"user","content":"hi"}],"stream":true}`
	resp, err := http.Post(proxy.URL+"/v1/chat/completions", "application/json", strings.NewReader(chatBody))
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		b, _ := io.ReadAll(resp.Body)
		t.Fatalf("status %d: %s", resp.StatusCode, string(b))
	}
	raw, _ := io.ReadAll(resp.Body)
	body := string(raw)
	if !strings.Contains(body, "data: ") || !strings.Contains(body, "hello") || !strings.Contains(body, "[DONE]") {
		t.Fatalf("stream body: %s", body)
	}
}

// TestStats verifies both streaming and non-streaming chat paths record token
// usage into the store and that GET /api/stats aggregates them correctly.
func TestStats(t *testing.T) {
	upstream, proxy, st := setupTestServer(t)
	defer upstream.Close()
	defer proxy.Close()

	creds := auth.Credentials{
		AccessKey: "AK", SecretKey: "SK", Token: "TOKEN",
		APIKey: "sk-test-456", ModelBaseURL: upstream.URL, Email: "dev@example.com",
	}
	if err := st.SaveCredentials(creds); err != nil {
		t.Fatal(err)
	}

	// non-streaming chat
	resp, err := http.Post(proxy.URL+"/v1/chat/completions", "application/json", strings.NewReader(`{"model":"GLM-4.7","messages":[{"role":"user","content":"hi"}],"stream":false}`))
	if err != nil {
		t.Fatal(err)
	}
	io.Copy(io.Discard, resp.Body)
	resp.Body.Close()

	// streaming chat
	resp, err = http.Post(proxy.URL+"/v1/chat/completions", "application/json", strings.NewReader(`{"model":"GLM-4.7","messages":[{"role":"user","content":"hi"}],"stream":true}`))
	if err != nil {
		t.Fatal(err)
	}
	io.Copy(io.Discard, resp.Body)
	resp.Body.Close()

	// query stats
	resp, err = http.Get(proxy.URL + "/api/stats")
	if err != nil {
		t.Fatal(err)
	}
	defer resp.Body.Close()
	var statsResp map[string]any
	if err := json.NewDecoder(resp.Body).Decode(&statsResp); err != nil {
		t.Fatal(err)
	}
	if enabled, _ := statsResp["enabled"].(bool); !enabled {
		t.Fatalf("stats not enabled: %v", statsResp)
	}
	s, _ := statsResp["stats"].(map[string]any)
	if s == nil {
		t.Fatalf("no stats object: %v", statsResp)
	}
	totals, _ := s["totals"].(map[string]any)
	if totals == nil {
		t.Fatalf("no totals: %v", s)
	}
	if got := totNum(totals["requests"]); got != 2 {
		t.Fatalf("requests = %v, want 2", totals["requests"])
	}
	if got := totNum(totals["prompt_tokens"]); got != 20 {
		t.Fatalf("prompt_tokens = %v, want 20", totals["prompt_tokens"])
	}
	if got := totNum(totals["completion_tokens"]); got != 40 {
		t.Fatalf("completion_tokens = %v, want 40", totals["completion_tokens"])
	}
	if got := totNum(totals["total_tokens"]); got != 60 {
		t.Fatalf("total_tokens = %v, want 60", totals["total_tokens"])
	}
	if got := totNum(totals["cached_tokens"]); got != 8 {
		t.Fatalf("cached_tokens = %v, want 8", totals["cached_tokens"])
	}
	if rate, _ := totals["cache_rate"].(float64); rate < 0.39 || rate > 0.41 {
		t.Fatalf("cache_rate = %v, want ~0.4", totals["cache_rate"])
	}
	perModel, _ := s["per_model"].([]any)
	if len(perModel) != 1 {
		t.Fatalf("per_model = %v, want 1 entry", perModel)
	}
}

// totNum extracts an int from a JSON-decoded numeric value (float64).
func totNum(v any) int {
	f, _ := v.(float64)
	return int(f)
}

func okValue(m map[string]any) bool {
	ok, _ := m["ok"].(bool)
	return ok
}

func TestMain(m *testing.M) {
	os.Exit(m.Run())
}

var _ = context.Background

const defaultTestPublicKey = `-----BEGIN PUBLIC KEY-----
MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAhxudxTewPgljUHEZHkusP7m3I+zA4/RGvuUMt6TtII/m4zwUOm/Y31zHBTmkCCt8k5vj9y+AmO0TsGmHooNQuMebakdmEWdcA5h7YAHHFbF2w5LcxIXjib08vgVpA+m3R5xPbLK+vfHe2aAX36b5nHReDNncY5vAl3U4CgIEBGPqyG67vJytRWqP+sfEdw5+m192Rf4SCGyiBzRmjiVlH3zeEBjdbOrkAnzKOVz6AHBl2q7LPLJKIzxjoAyhEp5qnDjHUFo5VZUgFwUOt83A/jbGMyzmjRoxBuvKcs9tBuorZyUwIsZN6E+rtQk2YqMPj4RkDsZ7LRmj6on8sN2rHQIDAQAB
-----END PUBLIC KEY-----`