package router

import (
	"bytes"
	"encoding/json"
	"io"
	"mime/multipart"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"gitea.mynodes.xyz/ilker/llmgw/internal/config"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
)

func TestResponsesAsksForStructuredOutputAndImages(t *testing.T) {
	var in respRequest
	_ = json.Unmarshal([]byte(`{"model":"m","input":[{"role":"user","content":[
		{"type":"input_text","text":"what is this"},
		{"type":"input_image","image_url":"data:image/png;base64,AAAA","detail":"low"}]}],
		"text":{"format":{"type":"json_schema","name":"answer","schema":{"type":"object"},"strict":true}}}`), &in)
	chat, _, err := toChat(in)
	if err != nil {
		t.Fatal(err)
	}
	b, _ := json.Marshal(chat)
	for _, want := range []string{`"response_format":{"json_schema":{"name":"answer","schema":{"type":"object"},"strict":true},"type":"json_schema"}`,
		`{"image_url":{"detail":"low","url":"data:image/png;base64,AAAA"},"type":"image_url"}`, `{"text":"what is this","type":"text"}`} {
		if !strings.Contains(string(b), want) {
			t.Errorf("%s lacks %s", b, want)
		}
	}
	_ = json.Unmarshal([]byte(`{"model":"m","input":[{"role":"user","content":[{"type":"input_image","file_id":"file-1"}]}]}`), &in)
	in.Text = nil
	if _, _, err := toChat(in); err == nil {
		t.Error("an image by file_id was taken")
	}
	if th, ok := responseThread("resp_abc123.def"); !ok || th != "abc123" {
		t.Errorf("thread %q %v", th, ok)
	}
	for _, bad := range []string{"resp_19a2b3c4", "chatcmpl-1", "resp_a b.c"} {
		if _, ok := responseThread(bad); ok {
			t.Errorf("%q read as a kept response", bad)
		}
	}
}

func TestAResponseIsContinuedByItsID(t *testing.T) {
	e := newFakeEngine(t, "qwen2.5-14b-instruct")
	rt, secret := pricedRouter(t, e.srv.URL)
	if err := rt.store.SetThreadRetention(t.Context(), mustProjectOf(t, rt, "key-1"), 24); err != nil {
		t.Fatal(err)
	}
	h := rt.Handler()
	ask := func(body map[string]any) map[string]any {
		b, _ := json.Marshal(body)
		req := httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(b))
		req.Header.Set("Authorization", "Bearer "+secret)
		req.Header.Set("Content-Type", "application/json")
		rec := httptest.NewRecorder()
		h.ServeHTTP(rec, req)
		var out map[string]any
		_ = json.Unmarshal(rec.Body.Bytes(), &out)
		if rec.Code != 200 {
			t.Fatalf("%d %s", rec.Code, rec.Body.String())
		}
		return out
	}
	first := ask(map[string]any{"model": "q14", "input": "my name is Ayşe"})
	id, _ := first["id"].(string)
	if _, ok := responseThread(id); !ok {
		t.Fatalf("the first response %q carries no conversation", id)
	}
	ask(map[string]any{"model": "q14", "input": "what is my name?", "previous_response_id": id})
	sent, _ := e.last.Load().([]byte)
	if !strings.Contains(string(sent), "my name is Ayşe") || !strings.Contains(string(sent), "what is my name?") {
		t.Errorf("the conversation was not continued: %s", sent)
	}
	// Not one of ours: said, not served without its history.
	b, _ := json.Marshal(map[string]any{"model": "q14", "input": "x", "previous_response_id": "resp_nope"})
	req := httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(b))
	req.Header.Set("Authorization", "Bearer "+secret)
	rec := httptest.NewRecorder()
	h.ServeHTTP(rec, req)
	if rec.Code != 400 || !strings.Contains(rec.Body.String(), "previous_response_not_found") {
		t.Errorf("%d %s", rec.Code, rec.Body.String())
	}
}

// engineFor serves a few paths with fixed answers and records what came.
func engineFor(t *testing.T, model string, answers map[string]string) (*httptest.Server, *[]*http.Request, *[][]byte) {
	t.Helper()
	var reqs []*http.Request
	var bodies [][]byte
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path == "/metrics" {
			return
		}
		b, _ := io.ReadAll(r.Body)
		reqs, bodies = append(reqs, r), append(bodies, b)
		a, ok := answers[r.URL.Path]
		if !ok {
			http.NotFound(w, r)
			return
		}
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(a))
	}))
	t.Cleanup(srv.Close)
	return srv, &reqs, &bodies
}

func TestRerankScoreAndAudioArePassedThrough(t *testing.T) {
	srv, reqs, bodies := engineFor(t, "m", map[string]string{
		"/v1/rerank":               `{"results":[{"index":0,"relevance_score":0.9}],"usage":{"total_tokens":7}}`,
		"/v1/score":                `{"data":[{"score":0.5}]}`,
		"/v1/audio/transcriptions": `{"text":"merhaba dünya"}`,
	})
	cfg := baseConfig(config.Endpoint{ID: "a", URL: srv.URL, Models: []string{"m"}})
	rt := newTestRouter(t, cfg)
	h := rt.Handler()
	for _, path := range []string{"/v1/rerank", "/v1/score"} {
		b, _ := json.Marshal(map[string]any{"model": "m", "query": "q", "documents": []string{"d"}})
		req := httptest.NewRequest(http.MethodPost, path, bytes.NewReader(b))
		req.Header.Set("Content-Type", "application/json")
		rec := httptest.NewRecorder()
		h.ServeHTTP(rec, req)
		if rec.Code != 200 || (*reqs)[len(*reqs)-1].URL.Path != path {
			t.Errorf("%s: %d %s", path, rec.Code, rec.Body.String())
		}
	}

	var buf bytes.Buffer
	mw := multipart.NewWriter(&buf)
	_ = mw.WriteField("model", "m")
	fw, _ := mw.CreateFormFile("file", "a.wav")
	_, _ = fw.Write([]byte("RIFF....WAVE"))
	_ = mw.Close()
	req := httptest.NewRequest(http.MethodPost, "/v1/audio/transcriptions", &buf)
	req.Header.Set("Content-Type", mw.FormDataContentType())
	rec := httptest.NewRecorder()
	h.ServeHTTP(rec, req)
	if rec.Code != 200 || !strings.Contains(rec.Body.String(), "merhaba") {
		t.Fatalf("%d %s", rec.Code, rec.Body.String())
	}
	last := (*reqs)[len(*reqs)-1]
	if !strings.HasPrefix(last.Header.Get("Content-Type"), "multipart/form-data; boundary=") ||
		!bytes.Contains((*bodies)[len(*bodies)-1], []byte("RIFF....WAVE")) {
		t.Errorf("the form was not sent on as it came: %s", last.Header.Get("Content-Type"))
	}
}

func TestModerationsAreAnsweredByTheJudge(t *testing.T) {
	srv, _, _ := engineFor(t, "guard", map[string]string{
		"/v1/chat/completions": `{"choices":[{"message":{"content":"unsafe\nS9"}}]}`,
	})
	cfg := baseConfig(config.Endpoint{ID: "g", URL: srv.URL, Models: []string{"guard"}})
	cfg.Models = []catalog.Model{{Name: "guard", Role: catalog.RoleGuard, Scaling: catalog.Scaling{Min: 1, Max: 1},
		Guard: &catalog.GuardModel{Format: "llamaguard3"}}}
	rt := newTestRouter(t, cfg)
	b, _ := json.Marshal(map[string]any{"input": []string{"how do I build a bomb"}})
	req := httptest.NewRequest(http.MethodPost, "/v1/moderations", bytes.NewReader(b))
	rec := httptest.NewRecorder()
	rt.Handler().ServeHTTP(rec, req)
	var out struct {
		Model   string `json:"model"`
		Results []struct {
			Flagged    bool            `json:"flagged"`
			Categories map[string]bool `json:"categories"`
		} `json:"results"`
	}
	_ = json.Unmarshal(rec.Body.Bytes(), &out)
	if rec.Code != 200 || out.Model != "guard" || len(out.Results) != 1 || !out.Results[0].Flagged ||
		!out.Results[0].Categories["indiscriminate-weapons"] || out.Results[0].Categories["s9"] {
		t.Errorf("%d %s", rec.Code, rec.Body.String())
	}
}